Update cross_encoder_reranking_train.py

cross_encoder_reranking_train.py

@@ -70,6 +70,29 @@ def process_single_patent(patent_dict):
         "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:
@@ -145,10 +168,8 @@ def extract_text(content_dict, text_type="full"):
         filtered_dict = process_single_patent(content_dict)
         all_text = []
         # Start with 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"])
+        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
@@ -162,6 +183,26 @@ def extract_text(content_dict, text_type="full"):
             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 paragraphs
+        for paragraph in filtered_dict["paragraphs"][:10]:
+            all_text.append(paragraph)
+        # Add features
+        for feature in filtered_dict["features"][:10]:
+            all_text.append(feature)
+        
+        return " ".join(all_text)
     
     
     return ""