Update utils/metadata.py

utils/metadata.py

@@ -2,64 +2,84 @@
 
 from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
 import re
+import dateparser
 
 # 🧠 Load advanced NER model
 model_name = "Jean-Baptiste/roberta-large-ner-english"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForTokenClassification.from_pretrained(model_name)
 
-# 🛠️ NER pipeline with aggregation
+# 🔧 Build NER pipeline with grouping
 ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
 
-
 def clean_text(text):
     """
-    Clean contract text for more accurate NER.
+    🧼 Clean contract text for better NER and regex performance.
     """
     return text.replace("\n", " ").replace("  ", " ").strip()
 
+def extract_effective_date(text):
+    """
+    📅 Extract natural language 'Effective Date' (e.g., 'as of August 28, 2025').
+    """
+    match = re.search(r"(?i)as of (.+?)(\.|,|\n)", text)
+    if match:
+        raw_date = match.group(1).strip()
+        parsed = dateparser.parse(raw_date)
+        if parsed:
+            return [parsed.strftime("%Y-%m-%d")]
+    return []
+
+def extract_parties(text):
+    """
+    🧾 Extract contracting parties using 'by and between X and Y'.
+    """
+    pattern = r"(?i)by and between[:\s\n]+(.+?)\s+and\s+(.+?)\s*(\(|\n|$)"
+    match = re.search(pattern, text, re.DOTALL)
+    if match:
+        return [match.group(1).strip(), match.group(2).strip()]
+    return []
 
 def extract_governing_law(text):
     """
-    Rule-based extractor for governing law clause.
+    ⚖️ Look for 'governed by the laws of XYZ'.
     """
     match = re.search(r"(?i)governed by the laws of ([\w\s,]+)", text)
-    return match.group(1).strip() if match else None
-
+    return [match.group(1).strip()] if match else []
 
 def extract_venue(text):
     """
-    Rule-based extractor for venue / jurisdiction clause.
+    🏛️ Look for venue in dispute clause like 'submitted to ... in XYZ'.
     """
     match = re.search(r"(?i)submitted to.*?in ([\w\s,]+)", text)
-    return match.group(1).strip() if match else None
-
+    return [match.group(1).strip()] if match else []
 
 def extract_metadata(text):
     """
-    Extract contract metadata using NER and rule-based matchers.
+    📦 Extract full structured metadata using hybrid rule-based + NER.
     """
     if not text.strip():
         return {"error": "No input provided."}
 
     text = clean_text(text)
-    max_chunk_length = 512  # safe for transformer models
+
+    # NER chunking
+    max_chunk_length = 512
     words = text.split()
     chunks = [" ".join(words[i:i + max_chunk_length]) for i in range(0, len(words), max_chunk_length)]
 
-    metadata = {
+    ner_metadata = {
         "EFFECTIVE_DATE": [],
         "PARTIES": [],
         "GOVERNING_LAW": [],
         "VENUE": []
     }
 
-    # NER label mapping → custom keys
     label_mapping = {
         "DATE": "EFFECTIVE_DATE",
         "PERSON": "PARTIES",
         "ORGANIZATION": "PARTIES",
-        "LOCATION": "GOVERNING_LAW"  # may include cities/states
+        "LOCATION": "GOVERNING_LAW"
     }
 
     for chunk in chunks:
@@ -68,15 +88,13 @@ def extract_metadata(text):
             label = ent["entity_group"]
             word = ent["word"]
             custom_label = label_mapping.get(label)
-            if custom_label and word not in metadata[custom_label]:
-                metadata[custom_label].append(word)
-
-    # 🧠 Override/inject rule-based values
-    law = extract_governing_law(text)
-    venue = extract_venue(text)
-    if law and law not in metadata["GOVERNING_LAW"]:
-        metadata["GOVERNING_LAW"].append(law)
-    if venue and venue not in metadata["VENUE"]:
-        metadata["VENUE"].append(venue)
-
-    return metadata
+            if custom_label and word not in ner_metadata[custom_label]:
+                ner_metadata[custom_label].append(word)
+
+    # 🧠 Replace/enhance with rule-based extraction
+    ner_metadata["PARTIES"] = extract_parties(text) or ner_metadata["PARTIES"]
+    ner_metadata["EFFECTIVE_DATE"] = extract_effective_date(text) or ner_metadata["EFFECTIVE_DATE"]
+    ner_metadata["GOVERNING_LAW"] = extract_governing_law(text) or ner_metadata["GOVERNING_LAW"]
+    ner_metadata["VENUE"] = extract_venue(text) or ner_metadata["VENUE"]
+
+    return ner_metadata