Upload evaluate.py

evaluate.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""
+Evaluate privacy-filter-enhanced on fax_number, credit_card_last4, company_contact_block.
+Compare against baseline openai/privacy-filter.
+"""
+
+import json
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+from seqeval.metrics import classification_report
+
+TEST_CASES = [
+    {
+        "text": "Please send the contract via fax to (555) 867-5309. Attn: Legal Dept.",
+        "expected": [("(555) 867-5309", "fax_number")]
+    },
+    {
+        "text": "Your Visa card ending in 4242 was charged $99.00 on 2024-01-15.",
+        "expected": [("4242", "credit_card_last4")]
+    },
+    {
+        "text": "Acme Corp\n123 Main St, Suite 100\nPhone: (555) 987-6543\nEmail: contact@acme.com",
+        "expected": [("Acme Corp\n123 Main St, Suite 100\nPhone: (555) 987-6543\nEmail: contact@acme.com", "company_contact_block")]
+    },
+    {
+        "text": "From: Alice Smith \u003calice@example.com\u003e\nFax: (212) 555-0199\nCard ending in 8765",
+        "expected": [("Alice Smith", "private_person"), ("alice@example.com", "private_email"),
+                     ("(212) 555-0199", "fax_number"), ("8765", "credit_card_last4")]
+    },
+]
+
+
+def predict(model_name, text):
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForTokenClassification.from_pretrained(model_name, trust_remote_code=True)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = model.to(device)
+    
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    preds = torch.argmax(logits, dim=2)[0].cpu().numpy()
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    labels = [model.config.id2label.get(p, "O") for p in preds]
+    
+    # Decode BIOES spans
+    spans = []
+    i = 0
+    while i < len(labels):
+        lab = labels[i]
+        if lab == "O" or lab.startswith("I-") or lab.startswith("E-"):
+            i += 1
+            continue
+        if lab.startswith("S-"):
+            cat = lab[2:]
+            spans.append((tokens[i], cat))
+            i += 1
+        elif lab.startswith("B-"):
+            cat = lab[2:]
+            j = i + 1
+            while j < len(labels) and labels[j] in (f"I-{cat}", f"E-{cat}"):
+                j += 1
+            span_tokens = tokens[i:j]
+            spans.append(("".join(span_tokens).replace("##", ""), cat))
+            i = j
+        else:
+            i += 1
+    return spans
+
+
+def main():
+    models = {
+        "baseline": "openai/privacy-filter",
+        "enhanced": "narcolepticchicken/privacy-filter-enhanced",
+    }
+    
+    for model_name, model_path in models.items():
+        print(f"\n{'='*60}")
+        print(f"Model: {model_name}")
+        print(f"{'='*60}")
+        for case in TEST_CASES:
+            text = case["text"]
+            expected = case["expected"]
+            try:
+                preds = predict(model_path, text)
+            except Exception as e:
+                preds = []
+                print(f"ERROR: {e}")
+            
+            print(f"\nText: {text[:80]}...")
+            print(f"Expected: {expected}")
+            print(f"Predicted: {preds}")
+            
+            # Simple match check
+            matched = 0
+            for exp_text, exp_label in expected:
+                found = any(exp_text in p[0] and p[1] == exp_label for p in preds)
+                if found:
+                    matched += 1
+            print(f"Matches: {matched}/{len(expected)}")
+
+
+if __name__ == "__main__":
+    main()