Update app.py

app.py

@@ -4,12 +4,15 @@ from fastapi.responses import JSONResponse
 from fastapi.staticfiles import StaticFiles
 from huggingface_hub import InferenceClient
 from sentence_transformers import SentenceTransformer
-
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 from fastapi import Request
 import requests
 import numpy as np
 import argparse
 import os
+from fastapi import HTTPException
+
 
 HOST = os.environ.get("API_URL", "0.0.0.0")
 PORT = os.environ.get("PORT", 7860)
@@ -30,8 +33,28 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+# Charge le tokenizer et le modèle
+tokenizer = AutoTokenizer.from_pretrained(
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    trust_remote_code=True
+)
+model = AutoModelForCausalLM.from_pretrained(
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    trust_remote_code=True,
+    torch_dtype=torch.float32,      # float32 sur CPU
+    low_cpu_mem_usage=True          # réduit l’empreinte mémoire
+)
+# Crée un pipeline "chat" (text-generation) préconfiguré
+chat_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=-1,                      # -1 = CPU
+    max_new_tokens=512,
+    temperature=0.7,
+    do_sample=True
+)
 
-app = FastAPI()
 embedder = SentenceTransformer('sentence-transformers/distiluse-base-multilingual-cased-v1')
 
 @app.post("/api/embed")
@@ -64,25 +87,24 @@ async def chat(request: Request):
     user_message = data.get("message", "").strip()
     if not user_message:
         raise HTTPException(status_code=400, detail="Le champ 'message' est requis.")
+    
+    # Construit le prompt
+    prompt = (
+        "Tu es un assistant médical spécialisé en schizophrénie.\n"
+        "Utilisateur : " + user_message + "\n"
+        "Assistant :"
+    )
 
     try:
-        # Appel au modèle en mode chat
-        completion = hf_client.chat.completions.create(
-            model="mistralai/Mistral-7B-Instruct-v0.3",
-            messages=[
-                {"role": "system", "content": "Tu es un assistant médical spécialisé en schizophrénie."},
-                {"role": "user",   "content": user_message}
-            ],
-            max_tokens=512,
-            temperature=0.7,
+        outputs = chat_pipeline(
+            prompt,
+            return_full_text=False
         )
-
-        bot_msg = completion.choices[0].message.content
+        bot_msg = outputs[0]["generated_text"].strip()
         return {"response": bot_msg}
 
     except Exception as e:
-        # En cas d'erreur d'inférence
-        raise HTTPException(status_code=502, detail=f"Erreur d'inférence HF : {e}")
+        raise HTTPException(status_code=502, detail=f"Erreur d’inférence locale : {e}")
 
 
 @app.get("/data")