Update app.py

app.py

@@ -1,229 +1,205 @@
 import gradio as gr
-import pandas as pd
 import numpy as np
 import h5py
+import faiss
 import json
-import os
-import tempfile
+from transformers import AutoTokenizer, AutoModel, AutoModelForMaskedLM
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 import re
-import time
-import logging
-from sentence_transformers import SentenceTransformer
+from collections import Counter
+import torch
 from nltk.corpus import stopwords
 from nltk.tokenize import word_tokenize
 import nltk
-import torch
-from sklearn.feature_extraction.text import CountVectorizer
 
-# Set up logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# Ensure you have downloaded the necessary NLTK data
+# Download necessary NLTK data
 nltk.download('stopwords', quiet=True)
 nltk.download('punkt', quiet=True)
 
-# Disable tokenizer parallelism warning
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
+# Load BERT model for lemmatization
+bert_lemma_model_name = "bert-base-uncased"
+bert_lemma_tokenizer = AutoTokenizer.from_pretrained(bert_lemma_model_name)
+bert_lemma_model = AutoModelForMaskedLM.from_pretrained(bert_lemma_model_name).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
-# Check for GPU availability
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Load BERT model for encoding search queries
+bert_encode_model_name = 'anferico/bert-for-patents'
+bert_encode_tokenizer = AutoTokenizer.from_pretrained(bert_encode_model_name)
+bert_encode_model = AutoModel.from_pretrained(bert_encode_model_name)
 
-# Load pre-trained model from Hugging Face
-logging.info("Loading SentenceTransformer model...")
-model = SentenceTransformer('anferico/bert-for-patents').to(device)
-logging.info("SentenceTransformer model loaded successfully.")
+def bert_lemmatize(text):
+    tokens = bert_lemma_tokenizer.tokenize(text)
+    input_ids = bert_lemma_tokenizer.convert_tokens_to_ids(tokens)
+    input_tensor = torch.tensor([input_ids]).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
+    with torch.no_grad():
+        outputs = bert_lemma_model(input_tensor)
+    predictions = outputs.logits.argmax(dim=-1)
+    lemmatized_tokens = bert_lemma_tokenizer.convert_ids_to_tokens(predictions[0])
+    return ' '.join([token for token in lemmatized_tokens if token not in ['[CLS]', '[SEP]', '[PAD]']])
 
-def preprocess_text(text):
-    # Remove "[EN]" label and claim numbers
-    text = re.sub(r'\[EN\]\s*', '', text)
-    text = re.sub(r'^\d+\.\s*', '', text, flags=re.MULTILINE)
-    
-    # Convert to lowercase while preserving acronyms and units
-    words = text.split()
-    text = ' '.join(word if word.isupper() or re.match(r'^\d+(\.\d+)?[a-zA-Z]+$', word) else word.lower() for word in words)
+def preprocess_query(text):
+    # Convert to lowercase
+    text = text.lower()
     
-    # Remove special characters except hyphens and periods in numbers
-    text = re.sub(r'[^\w\s\-.]', ' ', text)
-    text = re.sub(r'(?<!\d)\.(?!\d)', ' ', text)  # Remove periods not in numbers
+    # Remove any HTML tags (if present)
+    text = re.sub('<.*?>', '', text)
     
-    # Normalize spaces
-    text = re.sub(r'\s+', ' ', text).strip()
+    # Remove special characters, but keep hyphens, periods, and commas
+    text = re.sub(r'[^a-zA-Z0-9\s\-\.\,]', '', text)
     
     # Tokenize
     tokens = word_tokenize(text)
     
-    # Remove stopwords
+    # Remove stopwords, but keep all other words
     stop_words = set(stopwords.words('english'))
-    tokens = [word for word in tokens if word.lower() not in stop_words]
+    tokens = [word for word in tokens if word not in stop_words]
     
-    # Join tokens back into text
-    text = ' '.join(tokens)
+    # Join tokens back into a string
+    processed_text = ' '.join(tokens)
     
-    # Preserve numerical values with units
-    text = re.sub(r'(\d+(\.\d+)?)([a-zA-Z]+)', r'\1_\3', text)
+    # Apply BERT lemmatization
+    processed_text = bert_lemmatize(processed_text)
     
-    # Handle ranges and measurements
-    text = re.sub(r'(\d+(\.\d+)?)(\s*to\s*)(\d+(\.\d+)?)(\s*[a-zA-Z]+)', r'\1_to_\4_\6', text)
-    text = re.sub(r'between\s*(\d+(\.\d+)?)(\s*and\s*)(\d+(\.\d+)?)\s*([a-zA-Z]+)', r'between_\1_and_\4_\5', text)
-    
-    # Preserve chemical formulas
-    text = re.sub(r'\b([A-Z][a-z]?\d*)+\b', lambda m: m.group().replace(' ', ''), text)
-    
-    return text
+    return processed_text
 
-def filter_common_terms(texts, threshold=0.10):
-    vectorizer = CountVectorizer()
-    X = vectorizer.fit_transform(texts)
-    term_frequencies = np.sum(X.toarray(), axis=0)
-    document_frequencies = np.sum(X.toarray() > 0, axis=0)
-    num_documents = X.shape[0]
+def extract_key_features(text):
+    # For queries, we'll just preprocess and return all non-stopword terms
+    processed_text = preprocess_query(text)
     
-    common_terms = set()
-    removed_words = {}
-    for term, doc_freq in zip(vectorizer.get_feature_names_out(), document_frequencies):
-        if doc_freq / num_documents > threshold:
-            common_terms.add(term)
-            removed_words[term] = doc_freq
+    # Split the processed text into individual terms
+    features = processed_text.split()
     
-    filtered_texts = []
-    for text in texts:
-        filtered_text = ' '.join([word for word in text.split() if word not in common_terms])
-        filtered_texts.append(filtered_text)
+    # Remove duplicates while preserving order
+    features = list(dict.fromkeys(features))
     
-    return filtered_texts, removed_words
+    return features
 
-def encode_texts(texts, progress=gr.Progress(), batch_size=64):
-    embeddings = []
-    total_batches = len(texts) // batch_size + (1 if len(texts) % batch_size != 0 else 0)
-    
-    for i in range(0, len(texts), batch_size):
-        batch_texts = texts[i:i+batch_size]
-        batch_texts = [str(text) for text in batch_texts]
-        batch_embeddings = model.encode(batch_texts, show_progress_bar=True)
-        embeddings.extend(batch_embeddings)
-        progress((i // batch_size + 1) / total_batches, f"Processing batch {i // batch_size + 1}/{total_batches}")
-    
-    embeddings = np.array(embeddings)
-    embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
-    return embeddings
+def encode_texts(texts, max_length=512):
+    inputs = bert_encode_tokenizer(texts, padding=True, truncation=True, max_length=max_length, return_tensors='pt')
+    with torch.no_grad():
+        outputs = bert_encode_model(**inputs)
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+    return embeddings.numpy()
 
-def process_file(file, progress=gr.Progress()):
+def load_data():
     try:
-        start_time = time.time()
-        
-        # Read CSV file
-        df = pd.read_csv(file.name, encoding='utf-8')
-        logging.info(f"CSV file read successfully. Shape: {df.shape}")
-        
-        required_columns = ['Master Patent Number', 'Abstract', 'Claims']
-        missing_columns = [col for col in required_columns if col not in df.columns]
-        if missing_columns:
-            return None, None, None, f"Error: Missing columns: {', '.join(missing_columns)}"
-        
-        valid_texts = []
-        valid_patent_numbers = []
-        skipped_rows = []
-        error_rows = []
-        total_rows = len(df)
-        
-        for index, row in df.iterrows():
-            try:
-                progress((index + 1) / total_rows, f"Processing row {index + 1}/{total_rows}")
-                logging.info(f"Processing row {index + 1}/{total_rows}")
-                
-                abstract = row['Abstract'] if pd.notna(row['Abstract']) else ''
-                claims = row['Claims'] if pd.notna(row['Claims']) else ''
-                
-                if not abstract and not claims:
-                    skipped_rows.append(row['Master Patent Number'])
-                    continue
-                
-                # Preprocess the abstract and claims separately
-                preprocessed_abstract = preprocess_text(abstract)
-                preprocessed_claims = preprocess_text(claims)
-                
-                # Combine preprocessed abstract and claims
-                combined_text = preprocessed_abstract + ' ' + preprocessed_claims
-                
-                valid_texts.append(combined_text)
-                valid_patent_numbers.append(str(row['Master Patent Number']))
-                
-            except Exception as e:
-                error_message = f"Error processing row {index + 1}: {str(e)}"
-                logging.error(error_message)
-                error_rows.append((index, row['Master Patent Number'], error_message))
-                continue
-        
-        logging.info(f"Preprocessed abstracts and claims. Number of valid texts: {len(valid_texts)}")
-        
-        if skipped_rows:
-            logging.info(f"Skipped {len(skipped_rows)} rows due to missing abstract and claims.")
-        if error_rows:
-            logging.info(f"Encountered errors in {len(error_rows)} rows.")
+        with h5py.File('patent_embeddings.h5', 'r') as f:
+            embeddings = f['embeddings'][:]
+            patent_numbers = f['patent_numbers'][:]
         
-        # Filter out common terms
-        logging.info("Filtering common terms...")
-        filtered_texts, removed_words = filter_common_terms(valid_texts, threshold=0.10)
+        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'])
         
-        # Generate removed words file
-        removed_words_file = 'removed_words.txt'
-        with open(removed_words_file, 'w', encoding='utf-8') as f:
-            for word, count in sorted(removed_words.items(), key=lambda x: x[1], reverse=True):
-                f.write(f"{word}: {count}\n")
+        print(f"Embedding shape: {embeddings.shape}")
+        print(f"Number of patent numbers: {len(patent_numbers)}")
+        print(f"Number of metadata entries: {len(metadata)}")
         
-        logging.info("Encoding texts...")
-        embeddings = encode_texts(filtered_texts, progress)
-        logging.info("Texts encoded successfully.")
-        
-        # Save embeddings and metadata
-        embeddings_file = tempfile.NamedTemporaryFile(delete=False, suffix='.h5').name
-        with h5py.File(embeddings_file, 'w') as f:
-            f.create_dataset('embeddings', data=embeddings)
-            f.create_dataset('patent_numbers', data=valid_patent_numbers)
-        
-        metadata_file = tempfile.NamedTemporaryFile(delete=False, suffix='.jsonl').name
-        with open(metadata_file, 'w', encoding='utf-8') as f:
-            for index, (patent_number, text) in enumerate(zip(valid_patent_numbers, filtered_texts)):
-                json.dump({
-                    'index': index,
-                    'patent_number': patent_number,
-                    'text': text,
-                    'embedding_index': index
-                }, f, ensure_ascii=False)
-                f.write('\n')
-        
-        end_time = time.time()
-        total_time = end_time - start_time
-        logging.info(f"Processing completed in {total_time:.2f} seconds.")
-        
-        # Save error log
-        error_log_file = 'error_log.txt'
-        with open(error_log_file, 'w', encoding='utf-8') as f:
-            for row in error_rows:
-                f.write(f"Row {row[0]}, Patent {row[1]}: {row[2]}\n")
-        
-        return embeddings_file, metadata_file, removed_words_file, f"Processing complete. Encoded {len(filtered_texts)} patents. Skipped {len(skipped_rows)} patents due to missing data. Errors in {len(error_rows)} rows. See error_log.txt for details."
-    
+        return embeddings, patent_numbers, metadata, texts
+    except FileNotFoundError as e:
+        print(f"Error: Could not find file. {e}")
+        raise
     except Exception as e:
-        logging.error(f"An error occurred: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, None, None, f"An error occurred: {str(e)}"
+        print(f"An unexpected error occurred while loading data: {e}")
+        raise
+
+def compare_features(query_features, patent_features):
+    common_features = set(query_features) & set(patent_features)
+    similarity_score = len(common_features) / max(len(query_features), len(patent_features))
+    return common_features, similarity_score
+
+def hybrid_search(query, top_k=5):
+    print(f"Original query: {query}")
+    
+    processed_query = preprocess_query(query)
+    query_features = extract_key_features(processed_query)
+    
+    # Encode the processed query using the transformer model
+    query_embedding = encode_texts([processed_query])[0]
+    query_embedding = query_embedding / np.linalg.norm(query_embedding)
+    
+    # Perform semantic similarity search
+    semantic_distances, semantic_indices = index.search(np.array([query_embedding]).astype('float32'), top_k * 2)
+    
+    # Perform TF-IDF based search
+    query_tfidf = tfidf_vectorizer.transform([processed_query])
+    tfidf_similarities = cosine_similarity(query_tfidf, tfidf_matrix).flatten()
+    tfidf_indices = tfidf_similarities.argsort()[-top_k * 2:][::-1]
+    
+    # Combine and rank results
+    combined_results = {}
+    for i, idx in enumerate(semantic_indices[0]):
+        patent_number = patent_numbers[idx].decode('utf-8')
+        text = metadata[patent_number]['text']
+        patent_features = extract_key_features(text)
+        common_features, feature_similarity = compare_features(query_features, patent_features)
+        combined_results[patent_number] = {
+            'score': semantic_distances[0][i] * 1.0 + tfidf_similarities[idx] * 0.5 + feature_similarity,
+            'common_features': common_features,
+            'text': text
+        }
+    
+    for idx in tfidf_indices:
+        patent_number = patent_numbers[idx].decode('utf-8')
+        if patent_number not in combined_results:
+            text = metadata[patent_number]['text']
+            patent_features = extract_key_features(text)
+            common_features, feature_similarity = compare_features(query_features, patent_features)
+            combined_results[patent_number] = {
+                'score': tfidf_similarities[idx] * 1.0 + feature_similarity,
+                'common_features': common_features,
+                'text': text
+            }
+    
+    # Sort and get top results
+    top_results = sorted(combined_results.items(), key=lambda x: x[1]['score'], reverse=True)[:top_k]
+    
+    results = []
+    for patent_number, data in top_results:
+        result = f"Patent Number: {patent_number}\n"
+        result += f"Text: {data['text'][:200]}...\n"
+        result += f"Combined Score: {data['score']:.4f}\n"
+        result += f"Common Key Features: {', '.join(data['common_features'])}\n\n"
+        results.append(result)
+    
+    return "\n".join(results)
+
+# Load data and prepare the FAISS index
+embeddings, patent_numbers, metadata, texts = load_data()
+
+# Check if the embedding dimensions match
+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 = encode_texts(texts)
+    embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+
+# Normalize embeddings for cosine similarity
+embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
+
+# Create FAISS index for cosine similarity
+index = faiss.IndexFlatIP(embeddings.shape[1])
+index.add(embeddings)
+
+# Create TF-IDF vectorizer
+tfidf_vectorizer = TfidfVectorizer(stop_words='english')
+tfidf_matrix = tfidf_vectorizer.fit_transform(texts)
 
+# Create Gradio interface with additional input fields
 iface = gr.Interface(
-    fn=process_file,
-    inputs=gr.File(label="Upload a CSV file with patent data"),
-    outputs=[
-        gr.File(label="Patent Embeddings (HDF5)"),
-        gr.File(label="Patent Metadata (JSONL)"),
-        gr.File(label="Removed Words List (TXT)"),
-        gr.Textbox(label="Processing Status")
+    fn=hybrid_search,
+    inputs=[
+        gr.Textbox(lines=2, placeholder="Enter your patent query here..."),
+        gr.Slider(minimum=1, maximum=20, step=1, value=5, label="Top K Results"),
     ],
-    title="Patent Text Encoder",
-    description="Upload a CSV file containing patent data (must include 'Master Patent Number', 'Abstract', and 'Claims' columns). The app will generate embeddings and save them along with metadata as downloadable files.",
-    allow_flagging="never",
-    cache_examples=False,
+    outputs=gr.Textbox(lines=10, label="Search Results"),
+    title="Patent Similarity Search",
+    description="Enter a patent description to find similar patents based on key features."
 )
 
 if __name__ == "__main__":
-    iface.launch()
+    iface.launch()