Update app.py

app.py

@@ -5,7 +5,6 @@ import faiss
 import json
 from sentence_transformers import SentenceTransformer
 
-# Load embeddings and metadata
 def load_data():
     with h5py.File('patent_embeddings.h5', 'r') as f:
         embeddings = f['embeddings'][:]
@@ -17,6 +16,7 @@ def load_data():
             data = json.loads(line)
             metadata[data['patent_number']] = data
     
+    print(f"Embedding shape: {embeddings.shape}")
     return embeddings, patent_numbers, metadata
 
 embeddings, patent_numbers, metadata = load_data()
@@ -26,12 +26,30 @@ index = faiss.IndexFlatL2(embeddings.shape[1])
 index.add(embeddings)
 
 # Load BERT model for encoding search queries
-model = SentenceTransformer('all-MiniLM-L6-v2')
+embedding_dim = embeddings.shape[1]
+print(f"Embedding dimension: {embedding_dim}")
+
+if embedding_dim == 384:
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+elif embedding_dim == 768:
+    model = SentenceTransformer('all-mpnet-base-v2')
+else:
+    print(f"Unexpected embedding dimension: {embedding_dim}")
+    model = SentenceTransformer('all-MiniLM-L6-v2')  # Default to this model
 
 def search(query, top_k=5):
     # Encode the query
     query_embedding = model.encode([query])[0]
     
+    # Ensure the query embedding has the same dimension as the index
+    if query_embedding.shape[0] != index.d:
+        print(f"Query embedding dimension ({query_embedding.shape[0]}) does not match index dimension ({index.d})")
+        # Option 1: Pad or truncate the query embedding
+        if query_embedding.shape[0] < index.d:
+            query_embedding = np.pad(query_embedding, (0, index.d - query_embedding.shape[0]))
+        else:
+            query_embedding = query_embedding[:index.d]
+    
     # Perform similarity search
     distances, indices = index.search(np.array([query_embedding]), top_k)
     
@@ -55,4 +73,5 @@ iface = gr.Interface(
     description="Enter a query to find similar patents based on their embeddings."
 )
 
-iface.launch()
+if __name__ == "__main__":
+    iface.launch(share=True)