from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

app = FastAPI()

# Mount static files (HTML + CSS)
app.mount("/", StaticFiles(directory=".", html=True), name="static")

# Load model
model_name = "itsmeussa/AdabTranslate-Darija"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

@app.post("/translate")
async def translate(request: Request):
    data = await request.json()
    text = data.get("text", "")
    if not text.strip():
        return {"translation": "⚠️ Empty input."}

    inputs = tokenizer(text, return_tensors="pt")
    outputs = model.generate(**inputs, max_length=256)
    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"translation": translation}