from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, FileResponse
from transformers import AutoModelForCausalLM, AutoTokenizer
import uvicorn

app = FastAPI()

# Carica il modello Hugging Face
model_name = "microsoft/DialoGPT-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Servire il frontend statico
@app.get("/")
async def serve_index():
    return FileResponse("static/index.html")

# API per la chat
@app.post("/chat")
async def chat(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "")

    # Tokenizzazione e generazione della risposta
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=50)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return JSONResponse({"response": response})

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)