Update api.py

api.py

@@ -1,21 +1,19 @@
-from flask import Flask, request, jsonify
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 import os
 from huggingface_hub import login
+from flask import Flask, request, jsonify
 
-# Authenticate with your Hugging Face token
+# Authenticate with Hugging Face token from Secrets
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
-# Initialize Flask
-app = Flask(__name__)
+API_TOKEN = os.getenv("HF_API_TOKEN")  # Set this token in your Space Secrets
 
-# Model settings
+# Setup
 model_name = "cerebras/btlm-3b-8k-chat"
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 os.environ['HF_HOME'] = '/tmp/cache'
 
-# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
@@ -24,60 +22,63 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True
 )
 
-# Build text generation pipeline
 generator = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
     device_map="auto",
     torch_dtype=torch_dtype,
-    pad_token_id=tokenizer.eos_token_id
+    pad_token_id=tokenizer.eos_token_id,
+    trust_remote_code=True
 )
 
-# Route: /v1/chat/completions (OpenAI-compatible)
+app = Flask(__name__)
+
+@app.route("/")
+def home():
+    return "API is running"
+
 @app.route("/v1/chat/completions", methods=["POST"])
 def chat_completions():
-    try:
-        data = request.get_json()
-        messages = data.get("messages", [])
-        max_tokens = data.get("max_tokens", 256)
-        temperature = data.get("temperature", 0.7)
+    auth_header = request.headers.get("Authorization", "")
+    if not auth_header.startswith("Bearer ") or auth_header.split(" ")[1] != API_TOKEN:
+        return jsonify({"error": "Unauthorized"}), 401
 
-        # Extract user message from chat history
-        prompt = ""
-        for msg in messages:
-            role = msg.get("role", "")
-            content = msg.get("content", "")
-            if role and content:
-                prompt += f"{role.capitalize()}: {content}\n"
-        prompt += "Assistant:"
+    data = request.get_json()
+    messages = data.get("messages", [])
+    max_tokens = data.get("max_tokens", 256)
+    temperature = data.get("temperature", 0.7)
 
-        # Run generation
-        output = generator(
-            prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=0.9,
-            repetition_penalty=1.1,
-            do_sample=True
-        )
+    prompt = ""
+    for msg in messages:
+        role = msg.get("role", "")
+        content = msg.get("content", "")
+        if role and content:
+            prompt += f"{role.capitalize()}: {content}\n"
+    prompt += "Assistant:"
 
-        # Extract and format response
-        generated_text = output[0]['generated_text']
-        assistant_reply = generated_text.replace(prompt, "").strip()
+    output = generator(
+        prompt,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=0.9,
+        repetition_penalty=1.1,
+        do_sample=True
+    )
 
-        return jsonify({
-            "choices": [{
-                "message": {
-                    "role": "assistant",
-                    "content": assistant_reply
-                }
-            }]
-        })
+    generated_text = output[0]['generated_text']
+    assistant_reply = generated_text.replace(prompt, "").strip()
 
-    except Exception as e:
-        return jsonify({"error": str(e)}), 500
+    return jsonify({
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": assistant_reply
+            },
+            "finish_reason": "stop",
+            "index": 0
+        }]
+    })
 
-# Run Flask app
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=8080)