import os import json import torch import argparse import numpy as np from tqdm import tqdm from torch import Tensor import torch.nn.functional as F from transformers import AutoTokenizer, AutoModel from sklearn.cluster import AgglomerativeClustering from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity device = 'cuda' if torch.cuda.is_available() else 'cpu' # Load embedder once embedder = SentenceTransformer("all-MiniLM-L6-v2").to(device) # embedder = SentenceTransformer("AI-Growth-Lab/PatentSBERTa").to(device) def embed_text_list(texts): return embedder.encode(texts, convert_to_tensor=False, device=device) # def embed_text_list(texts): # # E5 models expect "query: " prefix for proper embedding behavior # formatted_texts = [f"query: {text}" for text in texts] # return embedder.encode(formatted_texts, convert_to_tensor=False, device=device) def rank_by_centrality(texts): embeddings = embed_text_list(texts) similarity_matrix = cosine_similarity(embeddings) centrality_scores = similarity_matrix.mean(axis=1) ranked = sorted(zip(texts, centrality_scores), key=lambda x: x[1], reverse=True) return [text for text, _ in ranked] def cluster_and_rank(texts, threshold=0.75): if len(texts) < 2: return texts embeddings = embed_text_list(texts) clustering = AgglomerativeClustering(n_clusters=None, distance_threshold=1-threshold, metric = "cosine", linkage='average') labels = clustering.fit_predict(embeddings) clustered_texts = {} for label, text in zip(labels, texts): clustered_texts.setdefault(label, []).append(text) representative_texts = [] for cluster_texts in clustered_texts.values(): ranked = rank_by_centrality(cluster_texts) representative_texts.append(ranked[0]) # Choose most central per cluster return representative_texts def process_single_patent(patent_dict): def filter_short_texts(texts, min_tokens=5): return [text for text in texts if len(text.split()) >= min_tokens] claims = filter_short_texts([v for k, v in patent_dict.items() if k.startswith("c-en")]) paragraphs = filter_short_texts([v for k, v in patent_dict.items() if k.startswith("p")]) features = filter_short_texts([v for k, v in patent_dict.get("features", {}).items()]) # Cluster & rank top_claims = cluster_and_rank(claims) top_paragraphs = cluster_and_rank(paragraphs) top_features = cluster_and_rank(features) return { "claims": rank_by_centrality(top_claims), "paragraphs": rank_by_centrality(top_paragraphs), "features": rank_by_centrality(top_features), } def process_single_patent2(patent_dict): def filter_short_texts(texts, min_tokens=5): return [text for text in texts if len(text.split()) >= min_tokens] # Filter short texts claims = filter_short_texts([v for k, v in patent_dict.items() if k.startswith("c-en")]) paragraphs = filter_short_texts([v for k, v in patent_dict.items() if k.startswith("p")]) features = filter_short_texts([v for k, v in patent_dict.get("features", {}).items()]) # Re-rank claims and features directly ranked_claims = rank_by_centrality(claims) ranked_features = rank_by_centrality(features) # Only filter (cluster + rank) for paragraphs filtered_paragraphs = cluster_and_rank(paragraphs) ranked_paragraphs = rank_by_centrality(filtered_paragraphs) return { "claims": ranked_claims, "paragraphs": ranked_paragraphs, "features": ranked_features, } def load_json_file(file_path): """Load JSON data from a file""" with open(file_path, 'r') as f: return json.load(f) def save_json_file(data, file_path): """Save data to a JSON file""" with open(file_path, 'w') as f: json.dump(data, f, indent=2) def load_content_data(file_path): """Load content data from a JSON file""" with open(file_path, 'r') as f: data = json.load(f) # Create a dictionary mapping FAN to Content content_dict = {item['FAN']: item['Content'] for item in data} return content_dict def extract_text(content_dict, text_type="full"): """Extract text from patent content based on text_type""" if text_type == "TA" or text_type == "title_abstract": # Extract title and abstract title = content_dict.get("title", "") abstract = content_dict.get("pa01", "") return f"{title} {abstract}".strip() elif text_type == "claims": # Extract all claims (keys starting with 'c') claims = [] for key, value in content_dict.items(): if key.startswith('c-'): claims.append(value) return " ".join(claims) elif text_type == "claimfeat": # Extract all claims (keys starting with 'c') content = [] for key, value in content_dict.items(): if key.startswith('c-'): content.append(value) if key == "features": content += list(content_dict[key].values()) return " ".join(content) elif text_type == "feat": # Extract all claims (keys starting with 'c') content = [] for key, value in content_dict.items(): if key == "features": content += list(content_dict[key].values()) return " ".join(content) elif text_type == "tac1": # Extract title, abstract, and first claim title = content_dict.get("title", "") abstract = content_dict.get("pa01", "") # Find the first claim safely first_claim = "" for key, value in content_dict.items(): if key.startswith('c-'): first_claim = value break return f"{title} {abstract} {first_claim}".strip() elif text_type == "description": # Extract all paragraphs (keys starting with 'p') paragraphs = [] for key, value in content_dict.items(): if key.startswith('p'): paragraphs.append(value) return " ".join(paragraphs) elif text_type == "full": # Extract everything all_text = [] # Start with title and abstract for better context at the beginning # if "title" in content_dict: # all_text.append(content_dict["title"]) # if "pa01" in content_dict: # all_text.append(content_dict["pa01"]) # Add claims and description for key, value in content_dict.items(): if key != "title" and key != "pa01": all_text.append(value) return " ".join(all_text) elif text_type == "smart": filtered_dict = process_single_patent(content_dict) all_text = [] # Start with abstract for better context at the beginning if "pa01" in content_dict: all_text.append(content_dict["pa01"]) # For claims, paragraphs and features, we take only the top-10 most relevant # Add claims # for claim in filtered_dict["claims"][:10]: # all_text.append(claim) # # Add features # for feature in filtered_dict["features"][:10]: # all_text.append(feature) # Add paragraphs for paragraph in filtered_dict["paragraphs"][:10]: all_text.append(paragraph) return " ".join(all_text) elif text_type == "smart2": filtered_dict = process_single_patent2(content_dict) all_text = [] # Start with abstract for better context at the beginning if "pa01" in content_dict: all_text.append(content_dict["pa01"]) # For claims, paragraphs and features, we take only the top-10 most relevant # Add claims for claim in filtered_dict["claims"][:10]: all_text.append(claim) # Add features for feature in filtered_dict["features"][:10]: all_text.append(feature) # # Add paragraphs # for paragraph in filtered_dict["paragraphs"][:10]: # all_text.append(paragraph) return " ".join(all_text) return "" def last_token_pool(last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor: """Extract the last token representations for pooling""" left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0]) if left_padding: return last_hidden_states[:, -1] else: sequence_lengths = attention_mask.sum(dim=1) - 1 batch_size = last_hidden_states.shape[0] return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths] def get_detailed_instruct(task_description: str, query: str) -> str: """Create an instruction-formatted query""" return f'Instruct: {task_description}\nQuery: {query}' def cross_encoder_reranking(query_text, doc_texts, model, tokenizer, batch_size=64, max_length=2048): """ Rerank document texts based on query text using cross-encoder model Parameters: query_text (str): The query text doc_texts (list): List of document texts model: The cross-encoder model tokenizer: The tokenizer for the model batch_size (int): Batch size for processing max_length (int): Maximum sequence length Returns: list: Indices of documents sorted by relevance score (descending) """ device = next(model.parameters()).device scores = [] # Format query with instruction task_description = 'Re-rank a set of retrieved patents based on their relevance to a given query patent. The task aims to refine the order of patents by evaluating their semantic similarity to the query patent, ensuring that the most relevant patents appear at the top of the list.' instructed_query = get_detailed_instruct(task_description, query_text) # Process in batches to avoid OOM for i in tqdm(range(0, len(doc_texts), batch_size), desc="Scoring documents", leave=False): batch_docs = doc_texts[i:i+batch_size] # Prepare input pairs for the batch input_texts = [instructed_query] + batch_docs # Tokenize with torch.no_grad(): batch_dict = tokenizer(input_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt').to(device) # Get embeddings outputs = model(**batch_dict) embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask']) # Normalize embeddings embeddings = F.normalize(embeddings, p=2, dim=1) # Calculate similarity scores between query and documents batch_scores = (embeddings[0].unsqueeze(0) @ embeddings[1:].T).squeeze(0) * 100 scores.extend(batch_scores.cpu().tolist()) # Create list of (index, score) tuples for sorting indexed_scores = list(enumerate(scores)) # Sort by score in descending order indexed_scores.sort(key=lambda x: x[1], reverse=True) # Return sorted indices return [idx for idx, _ in indexed_scores] def main(): base_directory = os.getcwd() base_directory += "/Patent_Retrieval" parser = argparse.ArgumentParser(description='Re-rank patents using cross-encoder scoring (training queries only)') parser.add_argument('--pre_ranking', type=str, default='shuffled_pre_ranking.json', help='Path to pre-ranking JSON file') parser.add_argument('--output', type=str, default='prediction2.json', help='Path to output re-ranked JSON file') parser.add_argument('--queries_content', type=str, default='./queries_content_with_features.json', help='Path to queries content JSON file') parser.add_argument('--documents_content', type=str, default='./documents_content_with_features.json', help='Path to documents content JSON file') # Change here for test or train parser.add_argument('--queries_list', type=str, default='test_queries.json', help='Path to training queries JSON file') parser.add_argument('--text_type', type=str, default='TA', choices=['TA', 'claims', 'description', 'full', 'tac1', 'smart', 'smart2', 'claimfeat', 'feat'], help='Type of text to use for scoring') parser.add_argument('--model_name', type=str, default='intfloat/e5-large-v2', help='Name of the cross-encoder model') parser.add_argument('--batch_size', type=int, default=4, help='Batch size for scoring') parser.add_argument('--max_length', type=int, default=512, help='Maximum sequence length') parser.add_argument('--device', type=str, default='cuda' if torch.cuda.is_available() else 'cpu', help='Device to use (cuda/cpu)') parser.add_argument('--base_dir', type=str, default=f'{base_directory}/datasets', help='Base directory for data files') args = parser.parse_args() # Ensure all paths are relative to base_dir if they're not absolute def get_full_path(path): if os.path.isabs(path): return path return os.path.join(args.base_dir, path) # Load training queries print(f"Loading training queries from {args.queries_list}...") queries_list = load_json_file(get_full_path(args.queries_list)) print(f"Loaded {len(queries_list)} training queries") # Load pre-ranking data print(f"Loading pre-ranking data from {args.pre_ranking}...") pre_ranking = load_json_file(get_full_path(args.pre_ranking)) # Filter pre-ranking to include only training queries pre_ranking = {fan: docs for fan, docs in pre_ranking.items() if fan in queries_list} print(f"Filtered pre-ranking to {len(pre_ranking)} training queries") # Load content data print(f"Loading query content from {args.queries_content}...") queries_content = load_content_data(get_full_path(args.queries_content)) print(f"Loading document content from {args.documents_content}...") documents_content = load_content_data(get_full_path(args.documents_content)) # Load model and tokenizer print(f"Loading model {args.model_name}...") tokenizer = AutoTokenizer.from_pretrained(args.model_name) model = AutoModel.from_pretrained(args.model_name).to(args.device) model.eval() # Process each query and re-rank its documents print("Starting re-ranking process for training queries...") re_ranked = {} missing_query_fans = [] missing_doc_fans = {} for query_fan, pre_ranked_docs in tqdm(pre_ranking.items(), desc="Processing queries"): # Check if query FAN exists in our content data if query_fan not in queries_content: missing_query_fans.append(query_fan) continue # Extract query text query_text = extract_text(queries_content[query_fan], args.text_type) if not query_text: missing_query_fans.append(query_fan) continue # Prepare document texts and keep track of their fans doc_texts = [] doc_fans = [] missing_docs_for_query = [] for doc_fan in pre_ranked_docs: if doc_fan not in documents_content: missing_docs_for_query.append(doc_fan) continue doc_text = extract_text(documents_content[doc_fan], args.text_type) if doc_text: doc_texts.append(doc_text) doc_fans.append(doc_fan) # Keep track of missing documents if missing_docs_for_query: missing_doc_fans[query_fan] = missing_docs_for_query # Skip if no valid documents if not doc_texts: re_ranked[query_fan] = [] continue # Re-rank documents print(f"\nRe-ranking {len(doc_texts)} documents for training query {query_fan}") # Print some of the original pre-ranking order for debugging print(f"Original pre-ranking (first 3): {doc_fans[:3]}") # Use cross-encoder model for reranking sorted_indices = cross_encoder_reranking( query_text, doc_texts, model, tokenizer, batch_size=args.batch_size, max_length=args.max_length ) re_ranked[query_fan] = [doc_fans[i] for i in sorted_indices] # Report any missing FANs if missing_query_fans: print(f"Warning: {len(missing_query_fans)} query FANs were not found in the content data") if missing_doc_fans: total_missing = sum(len(docs) for docs in missing_doc_fans.values()) print(f"Warning: {total_missing} document FANs were not found in the content data") # Save re-ranked results output_path = get_full_path(args.output) print(f"Saving re-ranked results to {output_path}...") save_json_file(re_ranked, output_path) print("Re-ranking complete!") print(f"Number of training queries processed: {len(re_ranked)}") # Optionally save the missing FANs information for debugging if missing_query_fans or missing_doc_fans: missing_info = { "missing_query_fans": missing_query_fans, "missing_doc_fans": missing_doc_fans } missing_info_path = f"{os.path.splitext(output_path)[0]}_missing_fans.json" save_json_file(missing_info, missing_info_path) print(f"Information about missing FANs saved to {missing_info_path}") if __name__ == "__main__": main()