Update app.py

app.py

@@ -44,11 +44,26 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 def get_prediction(prompt):
     st.write(f"Received prompt: {prompt}")  # Log the prompt received
+    
     # Tokenize the input prompt
     inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
     st.write(f"Tokenized input: {inputs}")  # Log the tokenized inputs
+
+    # Check if model is on the correct device
+    model.to(device)
+    
     # Generate output from the model
-    output = model.generate(inputs, max_new_tokens=150, temperature=0.7, top_p=0.95)
+    output = model.generate(
+        inputs, 
+        max_length=200,  # Set a reasonable max length for output
+        max_new_tokens=150,  # Limit output to avoid too long generations
+        temperature=0.7,  # Control randomness
+        top_p=0.95,  # Top-p sampling for diversity
+        do_sample=True,  # Enable sampling (for more diverse answers)
+        pad_token_id=tokenizer.eos_token_id,  # Ensure padding is handled
+        num_return_sequences=1  # Only generate 1 sequence
+    )
+
     st.write(f"Output: {output}")  # Log the raw output from the model
     # Decode the output to readable text
     decoded = tokenizer.decode(output[0], skip_special_tokens=True)
@@ -56,6 +71,7 @@ def get_prediction(prompt):
     return decoded.strip()
 
 
+
 # UI Header
 st.title("🧠 AnthroBot")
 st.write("Enter your anthropometric estimates to receive an interpreted summary inputs — manually or via CSV upload.")
@@ -113,4 +129,3 @@ with tab2:
                 csv_output = df.to_csv(index=False).encode("utf-8")
                 st.download_button("📤 Download Predictions", data=csv_output, file_name="predictions.csv")
 
-