Update app.py

app.py

@@ -1,34 +1,16 @@
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import pipeline
 
-app = Flask(__name__)
+app = FastAPI()
 
-# Charger le modèle depuis Hugging Face
-MODEL_NAME = "fatmata/psybot"  # Remplace avec le vrai nom de ton modèle
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
+# Charger le modèle hébergé sur Hugging Face
+chatbot = pipeline("text-generation", model="fatmata/psybot")  # Remplace par ton modèle
 
-@app.route("/chat", methods=["POST"])
-def chat():
-    data = request.json
-    user_input = data.get("message", "")
+class UserInput(BaseModel):
+    text: str
 
-    if not user_input:
-        return jsonify({"error": "Message vide"}), 400
-
-    # Génération de la réponse
-    prompt = f"<|startoftext|><|user|> {user_input} <|bot|>"
-    inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
-
-    with torch.no_grad():
-        output = model.generate(inputs, max_new_tokens=100, pad_token_id=tokenizer.eos_token_id)
-
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    if "<|bot|>" in response:
-        response = response.split("<|bot|>")[-1].strip()
-
-    return jsonify({"response": response})
-
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860)
+@app.post("/chatbot/")
+async def generate_response(user_input: UserInput):
+    response = chatbot(user_input.text, max_length=100, do_sample=True)[0]["generated_text"]
+    return {"response": response}