Spaces:
Running on Zero
Running on Zero
| import re | |
| import pandas as pd | |
| def postprocess(output: str) -> dict: | |
| """ | |
| Convert the model's output string to a dictionary. | |
| """ | |
| output = output.replace("<unk>", "#").removeprefix("<s>") | |
| pattern = ( | |
| r"(<LYSDAT>(?P<lysdat>.*?))?" | |
| r"(<VIGDAT>(?P<vigdat>.*?))?" | |
| r"((?P<kon1><MAN>|<KVINNA>)(?P<person1>.*?))" | |
| r"((?P<kon2><MAN>|<KVINNA>)(?P<person2>.*?))" | |
| r"(?=(<LYSDAT>|<VIGDAT>|</s>))" | |
| ) | |
| people = [] | |
| rad = 0 | |
| for m in re.finditer(pattern, output): | |
| res = m.groupdict() | |
| lysdat = res["lysdat"].strip() if res["lysdat"] else "" | |
| vigdat = res["vigdat"].strip() if res["vigdat"] else "" | |
| for i in range(2): | |
| person = f"person{i + 1}" | |
| kon = f"kon{i + 1}" | |
| p = {} | |
| for n in re.finditer(r"<(?P<key>.*?)>(?P<value>.*?)(?=<|$)", res[person]): | |
| key, value = n.groups() | |
| p["KON"] = res[kon].strip("<>").title() | |
| if key in ("OGIFT", "TROLOVAD", "ÄNKA_ÄNKLING", "SKILD"): | |
| p["CIV"] = key | |
| else: | |
| p[key] = value.strip() | |
| people.append({"LYSDAT": lysdat, "VIGDAT": vigdat, "RAD": rad} | p) | |
| rad += 1 | |
| return pd.DataFrame(people) | |