bhlewis's picture
Update app.py
74523b8 verified
raw
history blame
2.65 kB
import gradio as gr
import numpy as np
import h5py
import faiss
import json
from sentence_transformers import SentenceTransformer
def load_data():
with h5py.File('patent_embeddings.h5', 'r') as f:
embeddings = f['embeddings'][:]
patent_numbers = f['patent_numbers'][:]
metadata = {}
with open('patent_metadata.jsonl', 'r') as f:
for line in f:
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()
# Create FAISS index
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)
# Load BERT model for encoding search queries
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)
results = []
for i, idx in enumerate(indices[0]):
patent_number = patent_numbers[idx]
patent_data = metadata[patent_number]
result = f"Patent Number: {patent_number}\n"
result += f"Abstract: {patent_data['abstract'][:200]}...\n"
result += f"Similarity Score: {1 - distances[0][i]:.4f}\n\n"
results.append(result)
return "\n".join(results)
# Create Gradio interface
iface = gr.Interface(
fn=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."
)
if __name__ == "__main__":
iface.launch(share=True)