import os

# Imposta la cache dei modelli in una cartella scrivibile all'interno della home dell'utente
os.environ["HF_HOME"] = "/tmp/huggingface"

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

app = FastAPI()

# Crea la cartella di cache se non esiste
os.makedirs("/tmp/huggingface", exist_ok=True)

# Carica il modello Hugging Face
model_name = "microsoft/DialoGPT-medium"
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,
    pad_token_id=tokenizer.eos_token_id,  # Aggiunto per evitare warning
    attention_mask=inputs["attention_mask"]  # Aggiunto per maggiore stabilità
)
    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)