Update app.py

app.py

@@ -27,16 +27,26 @@ model = AutoModelForCausalLM.from_pretrained(
 ).to(device)
 
 # --- Define llm generation function ---
-def llm(prompt, max_new_tokens=400, do_sample=False):
+def llm(prompt, max_new_tokens=400, temperature=0.3, do_sample=True):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
+    generation_kwargs = {
+        "max_new_tokens": max_new_tokens,
+        "do_sample": do_sample,
+        "pad_token_id": tokenizer.eos_token_id,
+    }
+
+    # Only add temperature/top_p if sampling is enabled
+    if do_sample:
+        generation_kwargs.update({
+            "temperature": temperature,
+            "top_p": 0.95,
+            "top_k": 50
+        })
+
     output = model.generate(
         **inputs,
-        max_new_tokens=max_new_tokens,
-        do_sample=do_sample,
-        top_p=0.95 if do_sample else None,
-        top_k=50 if do_sample else None,
-        pad_token_id=tokenizer.eos_token_id
+        **generation_kwargs
     )
 
     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
@@ -45,6 +55,7 @@ def llm(prompt, max_new_tokens=400, do_sample=False):
 
 
 
+
 # Define all the screening questions
 questions = [
     # Generalized Anxiety & Somatic Concerns