Update app.py

app.py

@@ -2,17 +2,19 @@ from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 from huggingface_hub import login
+from flask import Flask, request, jsonify
 
-# Authentifiez-vous avec le jeton Hugging Face
+# Login to Hugging Face using secret token stored in Space secrets
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
-# Environnement de configuration et dtype pour la compatibilité CPU/GPU
+# Token authentication for requests
+API_TOKEN = os.getenv("HF_API_TOKEN")  # You set this in Space secrets
+
+# Set up model loading and pipeline
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 os.environ['HF_HOME'] = '/tmp/cache'
 
-# Charger le modèle et le tokenizer (en utilisant cerebras BTLM-3B-8K)
 model_name = "cerebras/btlm-3b-8k-chat"
-
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -21,40 +23,63 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
-# Créer un pipeline de génération de texte avec le pad_token_id requis pour ce modèle
 generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
     device_map="auto",
     torch_dtype=torch_dtype,
-    pad_token_id=tokenizer.eos_token_id,  # Important pour le modèle BTLM
+    pad_token_id=tokenizer.eos_token_id,
     trust_remote_code=True
 )
 
-def generate_chat_completion(message: str, history: list = None):
-    """
-    Si l'historique est fourni sous forme de liste de dictionnaires {'role': str, 'content': str},
-    reconstruit l'invite complète et renvoie l'historique mis à jour.
-    """
-    history = history or []
+# Flask app
+app = Flask(__name__)
+
+@app.route("/v1/chat/completions", methods=["POST"])
+def chat():
+    # Token auth: require Bearer token
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header.startswith("Bearer ") or auth_header.split(" ")[1] != API_TOKEN:
+        return jsonify({"error": "Unauthorized"}), 401
+
+    data = request.json
+    messages = data.get("messages", [])
+    max_tokens = data.get("max_tokens", 256)
+    temperature = data.get("temperature", 0.7)
+
+    # Build the prompt from chat history
     prompt = ""
-    for msg in history:
-        prompt += f"{msg['role'].capitalize()}: {msg['content']}\n"
-    prompt += f"Utilisateur: {message}\nAssistant:"
+    for msg in messages:
+        role = msg.get("role", "user").capitalize()
+        content = msg.get("content", "")
+        prompt += f"{role}: {content}\n"
+    prompt += "Assistant:"
 
+    # Generate response
     output = generator(
         prompt,
-        max_new_tokens=256,
-        temperature=0.7,  # Température légèrement inférieure pour des réponses plus cohérentes
+        max_new_tokens=max_tokens,
+        temperature=temperature,
         top_p=0.9,
         repetition_penalty=1.1,
         do_sample=True
     )
-    reply = output[0]['generated_text'].replace(prompt, "").strip()
+    reply = output[0]["generated_text"].replace(prompt, "").strip()
 
-    # Ajouter une nouvelle interaction à l'historique
-    history.append({"role": "user", "content": message})
-    history.append({"role": "assistant", "content": reply})
+    # Return response in OpenAI-style format
+    return jsonify({
+        "choices": [
+            {
+                "message": {
+                    "role": "assistant",
+                    "content": reply
+                },
+                "finish_reason": "stop",
+                "index": 0
+            }
+        ]
+    })
 
-    return history
+if __name__ == "__main__":
+    app.run(host="0.0.0.0", port=8081)