Update app.py

app.py

@@ -4,6 +4,8 @@ import h5py
 import faiss
 import json
 from sentence_transformers import SentenceTransformer
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 
 def load_data():
     try:
@@ -12,16 +14,18 @@ def load_data():
             patent_numbers = f['patent_numbers'][:]
         
         metadata = {}
+        texts = []
         with open('patent_metadata.jsonl', 'r') as f:
             for line in f:
                 data = json.loads(line)
                 metadata[data['patent_number']] = data
+                texts.append(data['text'])
         
         print(f"Embedding shape: {embeddings.shape}")
         print(f"Number of patent numbers: {len(patent_numbers)}")
         print(f"Number of metadata entries: {len(metadata)}")
         
-        return embeddings, patent_numbers, metadata
+        return embeddings, patent_numbers, metadata, texts
     except FileNotFoundError as e:
         print(f"Error: Could not find file. {e}")
         raise
@@ -29,7 +33,7 @@ def load_data():
         print(f"An unexpected error occurred while loading data: {e}")
         raise
 
-embeddings, patent_numbers, metadata = load_data()
+embeddings, patent_numbers, metadata, texts = load_data()
 
 # Normalize embeddings for cosine similarity
 embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
@@ -41,23 +45,43 @@ index.add(embeddings)
 # Load BERT model for encoding search queries
 model = SentenceTransformer('all-mpnet-base-v2')
 
-def search(query, top_k=5):
+# Create TF-IDF vectorizer
+tfidf_vectorizer = TfidfVectorizer(stop_words='english')
+tfidf_matrix = tfidf_vectorizer.fit_transform(texts)
+
+def hybrid_search(query, top_k=5):
     print(f"Searching for: {query}")
     
-    # Encode the query
+    # Encode the query using the transformer model
     query_embedding = model.encode([query])[0]
     query_embedding = query_embedding / np.linalg.norm(query_embedding)
     
-    print(f"Query embedding shape: {query_embedding.shape}")
+    # Perform semantic similarity search
+    semantic_distances, semantic_indices = index.search(np.array([query_embedding]), top_k * 2)
     
-    # Perform similarity search
-    distances, indices = index.search(np.array([query_embedding]), top_k)
+    # Perform TF-IDF based search
+    query_tfidf = tfidf_vectorizer.transform([query])
+    tfidf_similarities = cosine_similarity(query_tfidf, tfidf_matrix).flatten()
+    tfidf_indices = tfidf_similarities.argsort()[-top_k * 2:][::-1]
     
-    print(f"FAISS search results - Distances: {distances}, Indices: {indices}")
+    # Combine and rank results
+    combined_results = {}
+    for i, idx in enumerate(semantic_indices[0]):
+        patent_number = patent_numbers[idx].decode('utf-8')
+        combined_results[patent_number] = semantic_distances[0][i]
     
-    results = []
-    for i, idx in enumerate(indices[0]):
+    for idx in tfidf_indices:
         patent_number = patent_numbers[idx].decode('utf-8')
+        if patent_number in combined_results:
+            combined_results[patent_number] += tfidf_similarities[idx]
+        else:
+            combined_results[patent_number] = tfidf_similarities[idx]
+    
+    # Sort and get top results
+    top_results = sorted(combined_results.items(), key=lambda x: x[1], reverse=True)[:top_k]
+    
+    results = []
+    for patent_number, score in top_results:
         if patent_number not in metadata:
             print(f"Warning: Patent number {patent_number} not found in metadata")
             continue
@@ -65,20 +89,19 @@ def search(query, top_k=5):
         result = f"Patent Number: {patent_number}\n"
         text = patent_data.get('text', 'No text available')
         result += f"Text: {text[:200]}...\n"
-        result += f"Similarity Score: {distances[0][i]:.4f}\n\n"
+        result += f"Combined Score: {score:.4f}\n\n"
         results.append(result)
     
-    return "\n".join(results[:top_k])
+    return "\n".join(results)
 
 # Create Gradio interface
 iface = gr.Interface(
-    fn=search,
+    fn=hybrid_search,
     inputs=gr.Textbox(lines=2, placeholder="Enter your search query here..."),
     outputs=gr.Textbox(lines=10, label="Search Results"),
     title="Patent Similarity Search",
-    description="Enter a query to find similar patents based on their embeddings."
+    description="Enter a query to find similar patents based on their content."
 )
 
 if __name__ == "__main__":
     iface.launch()
-