from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

# Načítanie modelov
sk_to_en = pipeline("translation", model="Helsinki-NLP/opus-mt-sk-en")
en_to_sk = pipeline("translation", model="Helsinki-NLP/opus-mt-en-sk")

# Definuj vstupný formát
class TranslationRequest(BaseModel):
    data: list

@app.post("/predict/")
async def predict(req: TranslationRequest):
    try:
        text = req.data[0]
        direction = req.data[1]
        if direction == "sk-en":
            translation = sk_to_en(text)[0]["translation_text"]
        elif direction == "en-sk":
            translation = en_to_sk(text)[0]["translation_text"]
        else:
            return {"error": "Nepodporovaný smer prekladu."}
        return {"data": [translation]}
    except Exception as e:
        return {"error": str(e)}