from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline

app = FastAPI()

# Load model
tokenizer = AutoTokenizer.from_pretrained("ClergeF/ms-hardwick-translator", use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained("ClergeF/ms-hardwick-translator")
pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

class InputText(BaseModel):
    text: str

@app.post("/translate")
def translate(input: InputText):
    output = pipe(input.text, max_length=512, do_sample=False)
    return {"input": input.text, "output": output[0]["generated_text"]}