Upload inference_onnx.py with huggingface_hub

inference_onnx.py

@@ -29,10 +29,10 @@ class RgvedaEmbeddingONNXHybrid:
         print(f"Loading ONNX model: {model_path}")
         self.session = ort.InferenceSession(str(model_path))
         
-        # Load tokenizer (use the one from onnx-community for compatibility)
+        # Load tokenizer from local directory
         print("Loading tokenizer...")
         self.tokenizer = AutoTokenizer.from_pretrained(
-            "onnx-community/embeddinggemma-300m-ONNX"
+            str(self.model_dir)
         )
         
         # Load fine-tuned dense weights
@@ -47,6 +47,17 @@ class RgvedaEmbeddingONNXHybrid:
         print(f"  Dense1: {self.dense1_weight.shape}")
         print(f"  Dense2: {self.dense2_weight.shape}")
     
+    def mean_pooling(self, token_embeddings, attention_mask):
+        """Mean pooling with attention mask."""
+        input_mask_expanded = np.expand_dims(attention_mask, -1)
+        input_mask_expanded = np.broadcast_to(
+            input_mask_expanded, token_embeddings.shape
+        ).astype(np.float32)
+        
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.clip(np.sum(input_mask_expanded, axis=1), a_min=1e-9, a_max=None)
+        return sum_embeddings / sum_mask
+    
     def encode(self, texts, batch_size=32, show_progress=False):
         """
         Encode texts to embeddings using hybrid approach.
@@ -78,9 +89,9 @@ class RgvedaEmbeddingONNXHybrid:
             )
             
             # Run ONNX model
-            # The base model outputs: (last_hidden_state, sentence_embedding)
-            # where sentence_embedding already includes pooling + base dense layers
-            _, base_embedding = self.session.run(
+            # Get last_hidden_state (raw transformer output) not sentence_embedding
+            # sentence_embedding already has base dense layers which we don't want
+            last_hidden_state, _ = self.session.run(
                 None,
                 {
                     'input_ids': inputs['input_ids'].astype(np.int64),
@@ -88,12 +99,12 @@ class RgvedaEmbeddingONNXHybrid:
                 }
             )
             
-            # Apply fine-tuned dense layers
-            # Note: The base model already has dense layers, but we want to use
-            # the Rigveda-specific fine-tuned ones instead
+            # Do mean pooling ourselves (like the Ganaraj model does)
+            pooled = self.mean_pooling(last_hidden_state, inputs['attention_mask'])
             
+            # Now apply fine-tuned dense layers on the pooled output
             # Dense layer 1 (768 -> 3072)
-            dense1_out = base_embedding @ self.dense1_weight.T
+            dense1_out = pooled @ self.dense1_weight.T
             
             # Dense layer 2 (3072 -> 768)
             dense2_out = dense1_out @ self.dense2_weight.T