Update app.py

app.py

@@ -1,45 +1,46 @@
 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 torch
 import uvicorn
 
 app = FastAPI()
 
-# Crea la cartella di cache se non esiste
+# Imposta la cache per Hugging Face in una directory scrivibile
+os.environ["HF_HOME"] = "/tmp/huggingface"
 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)
+# Carica il modello DialoGPT
+model_name = "microsoft/DialoGPT-small"
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp/huggingface")
+model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="/tmp/huggingface")
 
-# 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})
+    # Tokenizzazione del prompt
+    input_ids = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
+
+    # Generazione della risposta
+    response_ids = model.generate(
+        input_ids,
+        max_length=100,
+        num_return_sequences=1,
+        pad_token_id=tokenizer.eos_token_id,
+        attention_mask=torch.ones(input_ids.shape, dtype=torch.long)  # Aggiunto per correggere l'errore
+    )
+
+    # Decodifica della risposta
+    response_text = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+
+    return JSONResponse({"response": response_text})
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)