Update api.py

api.py

@@ -4,22 +4,24 @@ import os
 from huggingface_hub import login
 from flask import Flask, request, jsonify
 
-# Authenticate with Hugging Face token from Secrets
+# Login with Hugging Face token
 login(os.getenv("HUGGINGFACEHUB_API_TOKEN"))
 
-API_TOKEN = os.getenv("HF_API_TOKEN")  # Set this token in your Space Secrets
+API_TOKEN = os.getenv("HF_API_TOKEN")
 
-# Setup
+# Model and loading config
 model_name = "cerebras/btlm-3b-8k-chat"
+revision = "main"  # Pin to specific model revision
 torch_dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
 os.environ['HF_HOME'] = '/tmp/cache'
 
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, revision=revision)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch_dtype,
     device_map="auto",
-    trust_remote_code=True
+    trust_remote_code=True,
+    revision=revision
 )
 
 generator = pipeline(
@@ -81,4 +83,4 @@ def chat_completions():
     })
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8080)
+    app.run(host="0.0.0.0", port=8080)