Update app.py

app.py

@@ -1,7 +1,46 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModel
+import torch
+import numpy as np
+import os
 
+# --- 1. Récupérer le token Hugging Face depuis variable d'environnement ---
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token is None:
+    raise ValueError("Tu dois définir la variable d'environnement HF_TOKEN avec ton token Hugging Face.")
 
+# --- 2. Charger SteelBERT pour embeddings ---
+steelbert_tokenizer = AutoTokenizer.from_pretrained(
+    "MGE-LLMs/SteelBERT", use_auth_token=hf_token
+)
+steelbert_model = AutoModel.from_pretrained(
+    "MGE-LLMs/SteelBERT", use_auth_token=hf_token
+).eval()
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+steelbert_model.to(device)
+
+def embed(text):
+    inputs = steelbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    with torch.no_grad():
+        outputs = steelbert_model(**inputs, output_hidden_states=True)
+    return outputs.hidden_states[-1][:,0,:].cpu().numpy()[0]
+
+# --- 3. Base documentaire (exemple, à remplacer par tes documents techniques) ---
+docs = {
+    "doc1": "L’acier X42 a une résistance à la traction de 415 MPa.",
+    "doc2": "L’acier inoxydable 304 est résistant à la corrosion et à l’oxydation."
+}
+doc_embeddings = {k: embed(v) for k,v in docs.items()}
+
+def search_best_doc(question):
+    q_emb = embed(question)
+    def cosine(a,b): return np.dot(a,b)/(np.linalg.norm(a)*np.linalg.norm(b))
+    best_doc = max(docs, key=lambda k: cosine(q_emb, doc_embeddings[k]))
+    return docs[best_doc]
+
+# --- 4. Fonction de réponse avec Mistral 7B Instruct ---
 def respond(
     message,
     history: list[dict[str, str]],
@@ -11,19 +50,21 @@ def respond(
     top_p,
     hf_token: gr.OAuthToken,
 ):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+    client = InferenceClient(token=hf_token.token, model="mistralai/Mistral-7B-Instruct-v0.2")
 
-    messages = [{"role": "system", "content": system_message}]
+    # Récupérer le contexte pertinent avec SteelBERT
+    best_doc = search_best_doc(message)
+    context = docs[best_doc]
 
+    # Construire le prompt
+    messages = [{"role": "system", "content": system_message}]
     messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
+    messages.append({
+        "role": "user",
+        "content": f"Question: {message}\nContexte: {context}\nRéponds clairement en français :"
+    })
 
     response = ""
-
     for message in client.chat_completion(
         messages,
         max_tokens=max_tokens,
@@ -31,32 +72,20 @@ def respond(
         temperature=temperature,
         top_p=top_p,
     ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+        if len(message.choices) and message.choices[0].delta.content:
+            token = message.choices[0].delta.content
+            response += token
+            yield response
 
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
+# --- 5. Interface Gradio ---
 chatbot = gr.ChatInterface(
     respond,
     type="messages",
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="Tu es un assistant spécialisé en métallurgie et en acier.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
 )
 
@@ -65,6 +94,5 @@ with gr.Blocks() as demo:
         gr.LoginButton()
     chatbot.render()
 
-
 if __name__ == "__main__":
     demo.launch()