Update models/loader.py

models/loader.py

@@ -29,11 +29,13 @@ MODEL_REGISTRY = {
 }
 _MODEL_CACHE = {}
 
-# Default generation config (can be overridden per call if needed)
+# Explicit generation config (avoids model-specific overrides)
 GENERATION_CONFIG = GenerationConfig(
     max_new_tokens=512,
+    do_sample=True,
     temperature=0.7,
     top_p=0.9,
+    top_k=50,
     repetition_penalty=1.1,
 )
 
@@ -48,14 +50,9 @@ def get_model_and_tokenizer(model_name):
 
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 
-        # Ensure a valid pad token exists
-        if tokenizer.pad_token is None:
-            if tokenizer.eos_token:
-                tokenizer.pad_token = tokenizer.eos_token
-            elif tokenizer.unk_token:
-                tokenizer.pad_token = tokenizer.unk_token
-            else:
-                tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
+        # Ensure a dedicated pad token exists (not EOS)
+        if tokenizer.pad_token is None or tokenizer.pad_token == tokenizer.eos_token:
+            tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
 
         # Load model with GPU/CPU awareness
         use_gpu = torch.cuda.is_available()
@@ -66,7 +63,7 @@ def get_model_and_tokenizer(model_name):
             trust_remote_code=True,
         )
 
-        # If new tokens were added, resize embeddings
+        # Resize embeddings if new tokens were added
         model.resize_token_embeddings(len(tokenizer))
 
         # Explicitly move to CPU if no GPU
@@ -87,11 +84,20 @@ def generate_with_model(agent_role, prompt, generation_config: GenerationConfig
 
     full_prompt = f"You are a helpful assistant. {ROLE_PROMPTS.get(agent_role, '')}\n\nUser prompt: {prompt}"
 
-    input_ids = tokenizer.encode(full_prompt, return_tensors="pt").to(model.device)
+    # Use tokenizer(...) to get both input_ids and attention_mask
+    inputs = tokenizer(
+        full_prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+    )
+    input_ids = inputs["input_ids"].to(model.device)
+    attention_mask = inputs["attention_mask"].to(model.device)
 
     with torch.no_grad():
         output = model.generate(
             input_ids,
+            attention_mask=attention_mask,   # ✅ ensures padding is ignored
             generation_config=generation_config,
             pad_token_id=tokenizer.pad_token_id,
         )