Update app.py

app.py

@@ -5,15 +5,15 @@ import torch
 import os
 
 # ---- LOAD LLM ----
-model_name = "mistralai/Mistral-7B-Instruct-v0.3"
+model_name = "meta-llama/Llama-3.2-3B-Instruct" 
 
 # Read token from environment variable (set in HF Space Secrets)
 hf_token = os.getenv("HF_TOKEN")
 
 tokenizer = AutoTokenizer.from_pretrained(
     model_name,
-    token=hf_token,
-    trust_remote_code=True
+    token=hf_token,         # Secure token
+    trust_remote_code=True  # Required for llama3 models
 )
 
 model = AutoModelForCausalLM.from_pretrained(
@@ -25,21 +25,28 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 
 # --- Define llm generation function ---
-def llm(prompt, max_new_tokens=500, temperature=0.3, do_sample=True):
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+def llm(prompt, max_new_tokens=1000, temperature=0.3, do_sample=True):
+    # Wrap the prompt into proper Llama 3 chat format
+    system_prompt = "[INST] " + prompt + " [/INST]"
+
+    inputs = tokenizer(system_prompt, return_tensors="pt").to(model.device)
+
     output = model.generate(
         **inputs,
         max_new_tokens=max_new_tokens,
         temperature=temperature,
-        do_sample=do_sample,  # <--- Use passed do_sample parameter
+        do_sample=do_sample,
         top_p=0.95,
         top_k=50,
         pad_token_id=tokenizer.eos_token_id
     )
+
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return [{"generated_text": generated_text}]
 
+    # Remove prompt part (optional cleanup)
+    generated_text = generated_text.replace(system_prompt, "").strip()
 
+    return [{"generated_text": generated_text}]
 
 
 # Define all the screening questions