Upload jepa_retrieval.py

jepa_retrieval.py

@@ -942,7 +942,7 @@ def setup_model_and_tokenizer(
     # Load model
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.bfloat16,
+        dtype=torch.bfloat16,
         device_map="auto",
         trust_remote_code=True,
         use_cache=False,
@@ -1074,15 +1074,67 @@ def run_inference(args):
     
     # Load model
     print("\n1. Loading model...")
-    tokenizer = AutoTokenizer.from_pretrained(args.model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        args.model_name,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-    )
     
-    if tokenizer.pad_token is None:
-        tokenizer.pad_token = tokenizer.eos_token
+    # Check if this is a LoRA/PEFT checkpoint (has adapter_config.json)
+    model_path = args.model_name
+    adapter_config_path = os.path.join(model_path, "adapter_config.json")
+    is_peft_checkpoint = os.path.exists(adapter_config_path)
+    
+    if is_peft_checkpoint:
+        print(f"Detected PEFT/LoRA checkpoint at {model_path}")
+        
+        # Read adapter config to get base model name
+        with open(adapter_config_path, 'r') as f:
+            adapter_config = json.load(f)
+        base_model_name = adapter_config.get("base_model_name_or_path", "meta-llama/Llama-3.2-1B-Instruct")
+        print(f"Base model: {base_model_name}")
+        
+        # Load tokenizer from checkpoint (has special tokens)
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        # Load base model
+        from peft import PeftModel
+        
+        base_model = AutoModelForCausalLM.from_pretrained(
+            base_model_name,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            local_files_only=False,
+        )
+        
+        # Resize embeddings to match tokenizer (with special tokens)
+        base_model.resize_token_embeddings(len(tokenizer))
+        
+        # Load PEFT adapter
+        model = PeftModel.from_pretrained(base_model, model_path)
+        model = model.merge_and_unload()  # Merge for faster inference
+        print("Loaded and merged PEFT adapter")
+        
+    else:
+        # Regular model (merged or base)
+        tokenizer = AutoTokenizer.from_pretrained(model_path)
+        
+        # Add special tokens (must match training)
+        special_tokens = ["[QUERY]", "[DOC]", "[SUPPORT]", "[ANSWER]"]
+        new_tokens = [t for t in special_tokens if t not in tokenizer.get_vocab()]
+        if new_tokens:
+            tokenizer.add_special_tokens({"additional_special_tokens": new_tokens})
+            print(f"Added {len(new_tokens)} special tokens: {new_tokens}")
+        
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+        
+        model = AutoModelForCausalLM.from_pretrained(
+            model_path,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+        
+        # Always resize embeddings to match tokenizer
+        model.resize_token_embeddings(len(tokenizer))
     
     # Load corpus
     print("\n2. Loading corpus...")