Fix attention mask warning by using tokenizer() and passing attention_mask to generate()

query_slm.py

@@ -15,7 +15,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 class LegalSLM:
     """Wrapper class for the fine-tuned legal SLM."""
     
-    def __init__(self, model_dir: str = "models/legal_slm"):
+    def __init__(self, model_dir: str = "."):
         """
         Initialize the Legal SLM.
         
@@ -34,6 +34,12 @@ class LegalSLM:
         self.model.to(self.device)
         self.model.eval()  # Set to evaluation mode
         
+        # Fix attention mask warning: GPT-2 uses same token for pad and eos
+        # Set pad_token_id explicitly and ensure it's handled correctly
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        
         print("Model loaded successfully!")
     
     def _sanitize_input(self, text: str, max_length: int = 1000) -> str:
@@ -99,20 +105,28 @@ class LegalSLM:
         # Build prompt
         prompt = f"Based on Australian legal documents, answer the following.\n\nQuestion: {sanitized_question}\nAnswer:"
         
-        # Tokenize prompt
-        inputs = self.tokenizer.encode(prompt, return_tensors='pt')
-        inputs = inputs.to(self.device)
+        # Tokenize prompt with attention mask to fix the warning
+        # Using tokenizer() instead of encode() to get attention_mask automatically
+        tokenized = self.tokenizer(
+            prompt,
+            return_tensors='pt',
+            padding=False,  # No padding needed for single input
+            truncation=False
+        )
+        input_ids = tokenized['input_ids'].to(self.device)
+        attention_mask = tokenized['attention_mask'].to(self.device)
         
         # Generate
         with torch.no_grad():
             outputs = self.model.generate(
-                inputs,
-                max_length=inputs.shape[1] + max_length,
+                input_ids,
+                attention_mask=attention_mask,  # Pass attention mask to fix warning
+                max_length=input_ids.shape[1] + max_length,
                 temperature=temperature,
                 top_p=top_p,
                 top_k=top_k,
                 do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id,
+                pad_token_id=self.tokenizer.pad_token_id,
                 eos_token_id=self.tokenizer.eos_token_id,
                 repetition_penalty=1.2,  # Reduce repetition
             )
@@ -181,8 +195,8 @@ def main():
     parser.add_argument(
         '--model-dir',
         type=str,
-        default='models/legal_slm',
-        help='Path to fine-tuned model directory'
+        default='.',
+        help='Path to fine-tuned model directory (default: current directory)'
     )
     parser.add_argument(
         '--question',