Update app.py

app.py

@@ -4,12 +4,12 @@ import h5py
 import faiss
 import json
 from transformers import AutoTokenizer, AutoModel
-from sentence_transformers import SentenceTransformer, models
 from sklearn.feature_extraction.text import TfidfVectorizer
 from sklearn.metrics.pairwise import cosine_similarity
 import re
 from collections import Counter
 import spacy
+import torch
 
 # Load Spacy model for advanced NLP
 try:
@@ -51,19 +51,24 @@ embeddings, patent_numbers, metadata, texts = load_data()
 try:
     tokenizer = AutoTokenizer.from_pretrained('anferico/bert-for-patents')
     bert_model = AutoModel.from_pretrained('anferico/bert-for-patents')
-    word_embedding_model = models.Transformer(model_name='anferico/bert-for-patents', tokenizer=tokenizer, auto_model=bert_model)
-    pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
-    model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
 except Exception as e:
     print(f"Error loading anferico/bert-for-patents: {e}")
     print("Falling back to a general-purpose model.")
-    model = SentenceTransformer('all-MiniLM-L6-v2')
+    tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+    bert_model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
+
+def encode_texts(texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
+    with torch.no_grad():
+        outputs = bert_model(**inputs)
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+    return embeddings.numpy()
 
 # Check if the embedding dimensions match
-if embeddings.shape[1] != model.get_sentence_embedding_dimension():
+if embeddings.shape[1] != encode_texts(["test"]).shape[1]:
     print("Embedding dimensions do not match. Rebuilding FAISS index.")
     # Rebuild embeddings using the new model
-    embeddings = np.array([model.encode(text) for text in texts])
+    embeddings = encode_texts(texts)
     embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
 
 # Normalize embeddings for cosine similarity
@@ -97,7 +102,7 @@ def hybrid_search(query, top_k=5):
     query_features = extract_key_features(query)
     
     # Encode the query using the transformer model
-    query_embedding = model.encode([query])[0]
+    query_embedding = encode_texts([query])[0]
     query_embedding = query_embedding / np.linalg.norm(query_embedding)
     
     # Perform semantic similarity search