Update model_loader.py

model_loader.py

@@ -5,14 +5,18 @@ from config import HF_TOKEN, MODEL_ID
 def load_model():
     try:
         print(f"🔄 Loading tokenizer and model: {MODEL_ID}")
-        
+
         # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_ID,
             token=HF_TOKEN,
-            trust_remote_code=True
+            trust_remote_code=True,
+            padding_side="left"  # For chat-style models
         )
 
+        # Set max length (MedGemma supports up to 8192 tokens)
+        tokenizer.model_max_length = 8192
+
         # Load model
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
@@ -25,12 +29,14 @@ def load_model():
 
         print("✅ Model loaded successfully.")
 
-        # Build pipeline
+        # Return generation pipeline with large max token output and context window
         return pipeline(
             "text-generation",
             model=model,
             tokenizer=tokenizer,
-            max_new_tokens=150,
+            max_new_tokens=1024,         # Max response length
+            truncation=True,             # Safely truncate input if needed
+            pad_token_id=tokenizer.eos_token_id,
             do_sample=True,
             temperature=0.7,
             top_p=0.9