Update space

app.py

@@ -75,6 +75,29 @@ MODEL_2_NAME = "sarvamai/sarvam-1"  # The base model on Hugging Face Hub
 # Load the tokenizer (same for both models)
 TOKENIZER_NAME = "sarvamai/sarvam-1"
 tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
+def fix_checkpoint(model_path):
+    """Fixes the model checkpoint by adjusting mismatched weight dimensions."""
+    checkpoint_file = os.path.join(model_path, "pytorch_model.bin")
+    fixed_checkpoint_file = os.path.join(model_path, "pytorch_model_fixed.bin")
+
+    if not os.path.exists(checkpoint_file):
+        raise FileNotFoundError(f"Checkpoint file not found at: {checkpoint_file}")
+
+    print("Loading checkpoint for fixing...")
+    checkpoint = torch.load(checkpoint_file, map_location="cpu")
+
+    # Adjust weights (truncate the last token if mismatch)
+    if "base_model.model.lm_head.base_layer.weight" in checkpoint:
+        checkpoint["base_model.model.lm_head.base_layer.weight"] = checkpoint["base_model.model.lm_head.base_layer.weight"][:-1]
+
+    if "base_model.model.lm_head.lora_B.default.weight" in checkpoint:
+        checkpoint["base_model.model.lm_head.lora_B.default.weight"] = checkpoint["base_model.model.lm_head.lora_B.default.weight"][:-1]
+
+    # Save the fixed checkpoint
+    print("Saving fixed checkpoint...")
+    torch.save(checkpoint, fixed_checkpoint_file)
+
+    return fixed_checkpoint_file  # Return the new file path
 
 # Function to load a model
 def load_model(model_choice):