Fix standalone model loading code - verified to work

README.md

@@ -100,20 +100,18 @@ class MultiModalEmbedder(nn.Module):
     
     def __init__(self):
         super().__init__()
-        # Text encoder
+        # Text encoder (384d, no projection needed)
         self.text_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
         self.text_encoder = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-        self.text_proj = nn.Linear(384, 384)
         
-        # Image encoder
+        # Image encoder (768d -> 384d projection)
         self.image_processor = SiglipProcessor.from_pretrained("google/siglip-base-patch16-512")
         self.image_encoder = SiglipModel.from_pretrained("google/siglip-base-patch16-512").vision_model
         self.image_proj = nn.Linear(768, 384)
         
-        # Audio encoder
+        # Audio encoder (384d, no projection needed)
         self.audio_processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny")
         self.audio_encoder = WhisperModel.from_pretrained("openai/whisper-tiny").encoder
-        self.audio_proj = nn.Linear(384, 384)
     
     def encode_text(self, texts):
         if isinstance(texts, str):
@@ -121,8 +119,7 @@ class MultiModalEmbedder(nn.Module):
         inputs = self.text_tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
         inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
         outputs = self.text_encoder(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)
-        embeddings = self.text_proj(embeddings)
+        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
         return F.normalize(embeddings, p=2, dim=-1)
     
     def encode_image(self, images):
@@ -130,16 +127,17 @@ class MultiModalEmbedder(nn.Module):
         inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
         outputs = self.image_encoder(**inputs)
         embeddings = outputs.pooler_output
-        embeddings = self.image_proj(embeddings)
+        embeddings = self.image_proj(embeddings)  # 768 -> 384
         return F.normalize(embeddings, p=2, dim=-1)
     
     def encode_audio(self, waveform):
-        # waveform: [batch, samples] at 16kHz
-        inputs = self.audio_processor(waveform.squeeze().numpy(), sampling_rate=16000, return_tensors="pt")
+        # waveform: numpy array or tensor at 16kHz
+        if isinstance(waveform, torch.Tensor):
+            waveform = waveform.squeeze().numpy()
+        inputs = self.audio_processor(waveform, sampling_rate=16000, return_tensors="pt")
         inputs = {k: v.to(next(self.parameters()).device) for k, v in inputs.items()}
         outputs = self.audio_encoder(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)
-        embeddings = self.audio_proj(embeddings)
+        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
         return F.normalize(embeddings, p=2, dim=-1)
 
 # Load model
@@ -150,15 +148,33 @@ checkpoint_path = hf_hub_download(
     repo_id="llm-semantic-router/multi-modal-embed-small",
     filename="model.pt"
 )
-state_dict = torch.load(checkpoint_path, map_location="cpu")
-
-# Map checkpoint keys to our model
-model.text_encoder.load_state_dict({k.replace("text_encoder.encoder.", ""): v for k, v in state_dict.items() if k.startswith("text_encoder.encoder.")})
-model.text_proj.load_state_dict({k.replace("text_encoder.projection.", ""): v for k, v in state_dict.items() if k.startswith("text_encoder.projection.")})
-model.image_encoder.load_state_dict({k.replace("image_encoder.vision_encoder.", ""): v for k, v in state_dict.items() if k.startswith("image_encoder.vision_encoder.")})
-model.image_proj.load_state_dict({k.replace("image_encoder.projection.", ""): v for k, v in state_dict.items() if k.startswith("image_encoder.projection.")})
-model.audio_encoder.load_state_dict({k.replace("audio_encoder.encoder.", ""): v for k, v in state_dict.items() if k.startswith("audio_encoder.encoder.")})
-model.audio_proj.load_state_dict({k.replace("audio_encoder.projection.", ""): v for k, v in state_dict.items() if k.startswith("audio_encoder.projection.")})
+state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+
+# Load text encoder weights
+model.text_encoder.load_state_dict({
+    k.replace("text_encoder.encoder.", ""): v 
+    for k, v in state_dict.items() 
+    if k.startswith("text_encoder.encoder.")
+})
+
+# Load image encoder and projection weights
+model.image_encoder.load_state_dict({
+    k.replace("image_encoder.vision_encoder.", ""): v 
+    for k, v in state_dict.items() 
+    if k.startswith("image_encoder.vision_encoder.")
+})
+model.image_proj.load_state_dict({
+    k.replace("image_encoder.projection.", ""): v 
+    for k, v in state_dict.items() 
+    if k.startswith("image_encoder.projection.")
+})
+
+# Load audio encoder weights
+model.audio_encoder.load_state_dict({
+    k.replace("audio_encoder.encoder.", ""): v 
+    for k, v in state_dict.items() 
+    if k.startswith("audio_encoder.encoder.")
+})
 
 model.eval()
 print("Model loaded successfully!")