from fastapi import FastAPI, Request
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline

app = FastAPI()

tokenizer = AutoTokenizer.from_pretrained("NlpHUST/vi-word-segmentation")
model = AutoModelForTokenClassification.from_pretrained("NlpHUST/vi-word-segmentation")
nlp = pipeline("token-classification", model=model, tokenizer=tokenizer)

class InputText(BaseModel):
    text: str

@app.post("/segment")
async def segment_text(payload: InputText):
    text = payload.text
    result = nlp(text)
    # Convert numpy values to native Python types
    processed_result = []
    for item in result:
        processed_item = {
            'entity': str(item['entity']),
            'score': float(item['score']),
            'word': str(item['word']),
            'start': int(item['start']),
            'end': int(item['end'])
        }
        processed_result.append(processed_item)
    return processed_result

@app.get("/")
def greet_json():
    return {"Hello": "World!"}