Update cross_encoder_reranking_train.py

cross_encoder_reranking_train.py

@@ -13,10 +13,13 @@ from sklearn.metrics.pairwise import cosine_similarity
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 # Load embedder once
-embedder = SentenceTransformer("sentence-transformers/all-mpnet-base-v2").to(device)
+# embedder = SentenceTransformer("all-MiniLM-L6-v2").to(device)
+embedder = SentenceTransformer("intfloat/e5-large-v2").to(device)
+
 
 def embed_text_list(texts):
-    return embedder.encode(texts, convert_to_tensor=False, device=device)
+    # return embedder.encode(texts, convert_to_tensor=False, device=device)
+    return embedder.encode(["query: your sentence here"], convert_to_tensor=False, device=device)
 
 def rank_by_centrality(texts):
     embeddings = embed_text_list(texts)
@@ -45,9 +48,12 @@ def cluster_and_rank(texts, threshold=0.75):
     return representative_texts
 
 def process_single_patent(patent_dict):
-    claims = [v for k, v in patent_dict.items() if k.startswith("c-en")]
-    paragraphs = [v for k, v in patent_dict.items() if k.startswith("p")]
-    features = [v for k, v in patent_dict.get("features", {}).items()]
+    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)
@@ -225,6 +231,7 @@ def cross_encoder_reranking(query_text, doc_texts, model, tokenizer, batch_size=
 
 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')
@@ -252,7 +259,7 @@ def main():
     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}/Patent_Retrieval/datasets',
+                        default=f'{base_directory}/datasets',
                         help='Base directory for data files')
     
     args = parser.parse_args()