Update app.py

app.py

@@ -15,6 +15,11 @@ logger = logging.getLogger(__name__)
 # Set cache directory for Hugging Face models
 os.environ["HF_HOME"] = "/tmp/huggingface"
 
+# Get Hugging Face token from environment variable (set in Spaces secrets)
+HF_TOKEN = os.getenv("HF_TOKEN")
+if not HF_TOKEN:
+    logger.warning("HF_TOKEN not set. Mistral model access may fail. Set it in Hugging Face Spaces secrets.")
+
 # Load dataset with error handling
 DATASET_PATH = os.path.join(os.getcwd(), "springer_papers_DL.json")
 try:
@@ -53,9 +58,9 @@ try:
     sci_bert_model.eval()
     logger.info("SciBERT loaded")
     
-    # Mistral-7B-Instruct for QA
-    mistral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", cache_dir="/tmp/huggingface")
-    mistral_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", cache_dir="/tmp/huggingface")
+    # Mistral-7B-Instruct for QA with token
+    mistral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", cache_dir="/tmp/huggingface", use_auth_token=HF_TOKEN)
+    mistral_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1", cache_dir="/tmp/huggingface", use_auth_token=HF_TOKEN)
     mistral_model.to(device)
     mistral_model.eval()
     logger.info("Mistral-7B-Instruct loaded")
@@ -160,14 +165,14 @@ def answer_question(paper, question, history):
         with torch.no_grad():
             outputs = mistral_model.generate(
                 inputs["input_ids"],
-                max_new_tokens=200,  # More tokens for detailed answers
+                max_new_tokens=200,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.9,
                 pad_token_id=mistral_tokenizer.eos_token_id
             )
         
-        # Decode and clean response
+        # Decode and clean response (preserve token structure)
         response = mistral_tokenizer.decode(outputs[0], skip_special_tokens=True)
         response = response[len(prompt):].strip()  # Remove prompt, including [INST] tags