Update app.py

app.py

@@ -1,32 +1,22 @@
-import os
-from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import pipeline
 
-# 🚀 Éviter le problème de cache dans Hugging Face Spaces
-os.environ["HF_HOME"] = "/app/huggingface_cache"
+app = FastAPI()
 
-# 📌 Charger le modèle depuis Hugging Face
-MODEL_NAME = "fatmata/psybot"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+# Charger le modèle Hugging Face
+model_name = "fatmata/psybot"
+generator = pipeline("text-generation", model=model_name)
 
-# 🌐 Créer l'application Flask
-app = Flask(__name__)
+# Définition du format des requêtes
+class TextRequest(BaseModel):
+    prompt: str
 
-@app.route('/chat', methods=['POST'])
-def chat():
-    """ Endpoint pour générer une réponse """
-    data = request.json
-    user_input = data.get("message", "")
+@app.get("/")
+def home():
+    return {"message": "Bienvenue sur l'API PsyBot !"}
 
-    if not user_input:
-        return jsonify({"error": "Aucun message fourni"}), 400
-
-    input_ids = tokenizer.encode(user_input, return_tensors="pt")
-    output = model.generate(input_ids, max_length=150, num_return_sequences=1)
-    response_text = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return jsonify({"response": response_text})
-
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=int(os.environ.get("PORT", 7860)), debug=False)
+@app.post("/generate/")
+def generate_text(request: TextRequest):
+    response = generator(request.prompt, max_length=100, do_sample=True)
+    return {"response": response[0]["generated_text"]}