from fastapi import FastAPI
from pydantic import BaseModel
from transformers import MarianTokenizer, MarianMTModel

tokenizer = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-de", clean_up_tokenization_spaces=True)
model = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-de")

def translate_text(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


app = FastAPI()

class TextIn(BaseModel):
    text: str

class TextOut(BaseModel):
    translation: str

@app.post("/translate", response_model=TextOut)
async def translate_api(data: TextIn):
    return {"translation": translate_text(data.text)}