conversion/fixtures.py: long inputs, <|mask|>, window sequences and span probabilities

conversion/fixtures.py

@@ -18,6 +18,7 @@ TEXTS = [
     "See https://example.com/a?b=c&d=e, then www.example.org.", "Ünïcödé façade — naïve café; Zürich, São Paulo.",
     "東京都千代田区1-1、山田太郎 (090-1234-5678)", "Привет, меня зовут Иван Петров.", "مرحبا اسمي أحمد", "tab\tseparated\tvalues\r\nCRLF",
     "IBAN GB82 WEST 1234 5698 7654 32, SSN 912-44-0387", "   leading and trailing   ", "emoji 👩‍👩‍👧 family 🇨🇦 flag", "",
+    "hello <|mask|> world", "x \t<|mask|>", "\n\n<|mask|> y", "a<|mask|>  b", "a" * 3000, "Ab" * 1500 + " end",
 ]
 tokens = []
 for s in TEXTS:
@@ -60,17 +61,23 @@ PII = [
     "Ünïcödé: Zoë Ågren lives at Kungsgatan 12, 111 43 Stockholm, phone +46 8 555 012 34.",
     "Contact: jose.garcia@example.org — José García, Calle Mayor 5, 28013 Madrid 🇪🇸",
 ]
+# long inputs: windowed PII (2 and 4 windows of the 1,024-token pass)
+PII += [" ".join(f"Record {i}: " + PII[i % 4] for i in range(n)) for n in (40, 90)]
 for s in PII:
     captured.clear(); sp = m.v.span_answer(s, {"type": "span", "criteria": json.load(open(sys.argv[1] + "/vela3_config.json"))["span"]["router_label_sets"]["pii"]})[0]
-    span_cases.append(dict(kind="pii", text=s, ids=captured[0], threshold=sp["threshold"],
-                           spans=[dict(label=x["label"], start=x["start"], end=x["end"]) for x in sp["spans"]]))
+    span_cases.append(dict(kind="pii", text=s, ids=list(captured), threshold=sp["threshold"],
+                           spans=[dict(label=x["label"], start=x["start"], end=x["end"], probability=x["probability"]) for x in sp["spans"]]))
 HALU = [("Adults may take up to 4 g of paracetamol in 24 hours.", "Adults may take up to 6 g of paracetamol in 24 hours."),
         ("The Eiffel Tower is 330 m tall and was finished in 1889.", "The Eiffel Tower, finished in 1899, is 330 m tall."),
         ("Refunds are issued within 5 business days to the original card.", "Refunds are issued within 5 business days to the original card."),
         ("The meeting moved from Tuesday to Thursday at 3 pm in room B.", "The meeting is now on Thursday at 4 pm in room C.")]
+# long inputs: a long context is cut (answer intact); a long answer is also read in windows
+LONG_CTX = " ".join(f"Note {i}: the warehouse in Lyon shipped {i * 7} boxes on day {i}." for i in range(120))
+HALU += [(LONG_CTX, "The warehouse in Lyon shipped 70 boxes on day 10 and 999 boxes on day 11."),
+         (LONG_CTX, " ".join(f"On day {i} it shipped {i * 7 + (1 if i % 9 == 0 else 0)} boxes." for i in range(150)))]
 for c, a in HALU:
     captured.clear(); sp = m.v.span_answer({"context": c, "answer": a}, {"type": "span", "criteria": ["unsupported"]})[0]
-    span_cases.append(dict(kind="unsupported", context=c, answer=a, ids=captured[0], threshold=sp["threshold"],
-                           spans=[dict(label=x["label"], start=x["start"], end=x["end"]) for x in sp["spans"]]))
+    span_cases.append(dict(kind="unsupported", context=c, answer=a, ids=list(captured), threshold=sp["threshold"],
+                           spans=[dict(label=x["label"], start=x["start"], end=x["end"], probability=x["probability"]) for x in sp["spans"]]))
 json.dump(dict(tokens=tokens, decisions=decisions, spans=span_cases), open(sys.argv[3], "w"), ensure_ascii=False)
 print(len(tokens), "token texts,", len(decisions), "decisions,", len(span_cases), "span calls ->", sys.argv[3])