Add attn mask

app.py

@@ -14,7 +14,6 @@ def load_model():
     
     return model, tokenizer
 
-# Function to generate SOAP notes
 def generate_soap_note(doctor_patient_conversation):
     if not doctor_patient_conversation.strip():
         return "Please enter a doctor-patient conversation."
@@ -26,11 +25,18 @@ Please generate a structured SOAP (Subjective, Objective, Assessment, Plan) note
 {doctor_patient_conversation}
 <|assistant|>"""
     
-    # Tokenize and generate
-    inputs = tokenizer(prompt, return_tensors="pt")
+    # Tokenize and generate with explicit padding settings
+    inputs = tokenizer(
+        prompt, 
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=tokenizer.model_max_length
+    )
     
     generate_ids = model.generate(
-        inputs.input_ids, 
+        inputs.input_ids,
+        attention_mask=inputs.attention_mask,  # Explicitly pass attention mask
         max_length=2048,
         num_beams=5, 
         no_repeat_ngram_size=2,