| import pandas as pd | |
| import os | |
| import re | |
| import csv | |
| def extract_paren(annotation): | |
| ents = [] | |
| for i in range(len(annotation)): | |
| if annotation[i] == "[": | |
| ent = "[" | |
| open_paren = 0 | |
| for j in range(i+1, len(annotation)): | |
| if annotation[j] == "[": | |
| open_paren += 1 | |
| elif annotation[j] == "]": | |
| if open_paren > 0: | |
| open_paren -= 1 | |
| ent = ent[:len(ent)-3] | |
| else: | |
| ent += "]" | |
| digit = re.search(r": [0-9]{1,3}", ent) | |
| if digit: | |
| matches = re.findall(r": [0-9]{1,3}", annotation[:i]) | |
| str_index = annotation[:i].count(" ") - len(matches) | |
| ent += "|" + str(str_index) | |
| ents.append(ent) | |
| break | |
| else: | |
| ent += annotation[j] | |
| return ents | |
| def create_clusters(ents): | |
| clusters = {} | |
| for e in ents: | |
| digit_ann = re.search(r": [0-9]{1,3}", e) | |
| if digit_ann: | |
| clean_e = e.replace("[", "").replace("]", "").replace(digit_ann.group(), "") | |
| digit = re.search(r"[0-9]{1,3}", digit_ann.group()) | |
| digit = int(digit.group()) | |
| if digit not in clusters: | |
| clusters[digit] = [] | |
| clusters[digit].append(clean_e) | |
| else: | |
| print("OH NO:", e) | |
| print() | |
| return clusters | |
| headers = ["input", "model_output", "model_output_clusters"] | |
| df = pd.read_csv("results.csv") | |
| rows = [] | |
| for index, row in df.iterrows(): | |
| annotation = row["model_output"] | |
| if isinstance(annotation, str): | |
| ann_ents = extract_paren(annotation) | |
| ann_clusters = {} | |
| if ann_ents: | |
| ann_clusters = create_clusters(ann_ents) | |
| else: | |
| ann_clusters = {} | |
| new_row = [row["input"], annotation, str(ann_clusters)] | |
| rows.append(new_row) | |
| f = open("cluster_results.csv", "w") | |
| writer = csv.writer(f) | |
| writer.writerow(headers) | |
| writer.writerows(rows) | |
| f.close() |