Make README fully standalone - no external repo needed

README.md

@@ -85,23 +85,8 @@ pip install torch transformers pillow safetensors
 ### Load Model
 
 Two checkpoint formats are available:
-- `model.pt` (932 MB) - PyTorch format, smaller due to shared tensors
-- `model.safetensors` (1.35 GB) - SafeTensors format, recommended for production
-
-**Option 1: Using the source repository (full features)**
-
-```bash
-git clone https://github.com/semantic-router/2DMSE-Multimodal-Embedder.git
-cd 2DMSE-Multimodal-Embedder
-pip install -e .
-```
-
-```python
-from src.models import MultimodalEmbedder
-model = MultimodalEmbedder.from_pretrained("llm-semantic-router/multi-modal-embed-small")
-```
-
-**Option 2: Standalone with transformers (no repo needed)**
+- `model.pt` (932 MB) - PyTorch format
+- `model.safetensors` (1.35 GB) - SafeTensors format
 
 ```python
 import torch
@@ -110,53 +95,75 @@ import torch.nn.functional as F
 from transformers import AutoModel, AutoTokenizer, SiglipModel, SiglipProcessor, WhisperModel, WhisperFeatureExtractor
 from huggingface_hub import hf_hub_download
 
-# Download weights
+class MultiModalEmbedder(nn.Module):
+    """Standalone multimodal embedder - no external dependencies."""
+    
+    def __init__(self):
+        super().__init__()
+        # Text encoder
+        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
+        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
+        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):
+            texts = [texts]
+        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)
+        return F.normalize(embeddings, p=2, dim=-1)
+    
+    def encode_image(self, images):
+        inputs = self.image_processor(images=images, return_tensors="pt")
+        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)
+        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")
+        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)
+        return F.normalize(embeddings, p=2, dim=-1)
+
+# Load model
+model = MultiModalEmbedder()
+
+# Download and load trained weights
 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")
 
-# Load text encoder
-text_tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-text_encoder = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-
-# Load image encoder  
-image_processor = SiglipProcessor.from_pretrained("google/siglip-base-patch16-512")
-image_encoder = SiglipModel.from_pretrained("google/siglip-base-patch16-512").vision_model
-
-# Load audio encoder
-audio_processor = WhisperFeatureExtractor.from_pretrained("openai/whisper-tiny")
-audio_encoder = WhisperModel.from_pretrained("openai/whisper-tiny").encoder
-
-# Load trained projection weights from checkpoint
-# Text projection: state_dict keys starting with "text_encoder.projection"
-# Image projection: state_dict keys starting with "image_encoder.projection"
-# Audio projection: state_dict keys starting with "audio_encoder.projection"
-
-def encode_text(texts, tokenizer=text_tokenizer, encoder=text_encoder):
-    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
-    with torch.no_grad():
-        outputs = encoder(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)  # Mean pooling
-        embeddings = F.normalize(embeddings, p=2, dim=-1)
-    return embeddings
-
-def encode_image(images, processor=image_processor, encoder=image_encoder):
-    inputs = processor(images=images, return_tensors="pt")
-    with torch.no_grad():
-        outputs = encoder(inputs.pixel_values)
-        embeddings = outputs.pooler_output
-        embeddings = F.normalize(embeddings, p=2, dim=-1)
-    return embeddings
+# 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.")})
 
-# Example usage
-text_emb = encode_text(["A photo of a cat"])
-print(f"Text embedding shape: {text_emb.shape}")
+model.eval()
+print("Model loaded successfully!")
 ```
 
-> **Note**: Option 2 loads the base encoders but not the trained projection layers. For full model with trained weights, use Option 1.
-
 ### Text Embedding
 
 ```python