Spaces:
Runtime error
Runtime error
import json | |
import copy | |
import torch | |
import torch.nn.functional as F | |
import numpy as np | |
import faiss | |
from tqdm import tqdm | |
from torch.utils.data import Dataset, DataLoader | |
from torch.cuda.amp import autocast | |
court_text_splitter = "Весь текст судебного документа: " | |
class FaissDocsDataset(Dataset): | |
def __init__(self, data): | |
self.data = data | |
def __len__(self): | |
return len(self.data) | |
def __getitem__(self, idx): | |
return self.data[idx] | |
def preprocess_inputs(inputs, device): | |
return {k: v[:, 0, :].to(device) for k, v in inputs.items()} | |
def get_subsets_for_db(subsets, data_ids, all_docs): | |
subsets = [data_ids[ss_name] for ss_name in subsets] | |
subsets = [x for ss in subsets for x in ss] | |
all_docs_db = {k: v for k, v in all_docs.items() | |
if v['id'] in subsets} | |
unique_refs = set([ref for doc in all_docs_db.values() | |
for ref, text in doc['added_refs'].items()]) | |
db_data = {ref: text for doc in all_docs_db.values() | |
for ref, text in doc['added_refs'].items() if ref in unique_refs} | |
return db_data | |
def get_subsets_for_qa(subsets, data_ids, all_docs): | |
subsets = [data_ids[ss_name] for ss_name in subsets] | |
subsets = [x for ss in subsets for x in ss] | |
all_docs_qa = {k: v for k, v in all_docs.items() | |
if v['id'] in subsets} | |
return all_docs_qa | |
def filter_db_data_types(text_parts, db_data_in): | |
filtered_db_data = {} | |
db_data = copy.deepcopy(db_data_in) | |
for ref, text in db_data.items(): | |
if any([True for x in text_parts if x in ref]): | |
filtered_db_data[ref] = text | |
return filtered_db_data | |
def filter_qa_data_types(text_parts, all_docs_in): | |
filtered_all_docs = {} | |
all_docs = copy.deepcopy(all_docs_in) | |
for doc_key, doc in all_docs.items(): | |
if not len(doc['added_refs']): | |
filtered_all_docs[doc_key] = doc | |
continue | |
filtered_refs = {} | |
for ref, text in doc['added_refs'].items(): | |
if any([True for x in text_parts if x in ref]): | |
filtered_refs[ref] = text | |
filtered_all_docs[doc_key] = doc | |
filtered_all_docs[doc_key]['added_refs'] = filtered_refs | |
return filtered_all_docs | |
def db_tokenization(filtered_db_data, tokenizer, max_len=510): | |
index_keys = {} | |
index_toks = {} | |
for key_idx, (ref, text) in enumerate(tqdm(filtered_db_data.items(), | |
desc="Tokenizing DB refs")): | |
index_keys[key_idx] = ref | |
text = "passage: " + text | |
index_toks[key_idx] = tokenizer(text, return_tensors="pt", | |
padding='max_length', truncation=True, | |
max_length=max_len) | |
return index_keys, index_toks | |
def qa_tokenization(all_docs_qa, tokenizer, max_len=510): | |
ss_docs = [] | |
for doc in tqdm(all_docs_qa.values(), desc="Tokenizing QA docs"): | |
text = doc['title'] + '\n' + doc['question'] | |
text = "query: " + text | |
text = tokenizer(text, return_tensors="pt", | |
padding='max_length', truncation=True, | |
max_length=max_len) | |
ss_docs.append([text, list(doc['added_refs'].keys())]) | |
val_questions = [x[0] for x in ss_docs] | |
val_refs = {idx: x[1] for idx, x in enumerate(ss_docs)} | |
return val_questions, val_refs | |
def query_tokenization(text, tokenizer, max_len=510): | |
text = "query: " + text | |
text = tokenizer(text, return_tensors="pt", | |
padding='max_length', truncation=True, | |
max_length=max_len) | |
return text | |
def query_embed_extraction(tokens, model, do_normalization=True): | |
model.eval() | |
device = model.device | |
with torch.no_grad(): | |
with autocast(): | |
inputs = {k: v[:, :].to(device) for k, v in tokens.items()} | |
outputs = model(**inputs) | |
embedding = outputs.last_hidden_state[:, 0].cpu() | |
if do_normalization: | |
embedding = F.normalize(embedding, dim=-1) | |
return embedding.numpy() | |
def extract_text_embeddings(index_toks, val_questions, model, | |
do_normalization=True, faiss_batch_size=16): | |
faiss_dataset = FaissDocsDataset(list(index_toks.values())) | |
db_data_loader = DataLoader(faiss_dataset, batch_size=faiss_batch_size) | |
ss_val_dataset = FaissDocsDataset(val_questions) | |
qu_data_loader = DataLoader(ss_val_dataset, batch_size=faiss_batch_size) | |
model.eval() | |
device = model.device | |
docs_embeds = [] | |
questions_embeds = [] | |
with torch.no_grad(): | |
for batch in tqdm(db_data_loader, desc="db_embeds_extraction"): | |
with autocast(): | |
outputs = model(**preprocess_inputs(batch, device)) | |
docs_embeds.extend(outputs.last_hidden_state[:, 0].cpu()) | |
for batch in tqdm(qu_data_loader, desc="qu_embeds_extraction"): | |
with autocast(): | |
outputs = model(**preprocess_inputs(batch, device)) | |
questions_embeds.extend(outputs.last_hidden_state[:, 0].cpu()) | |
docs_embeds_faiss = [torch.unsqueeze(x, 0) for x in docs_embeds] | |
docs_embeds_faiss = torch.cat(docs_embeds_faiss) | |
questions_embeds_faiss = [torch.unsqueeze(x, 0) for x in questions_embeds] | |
questions_embeds_faiss = torch.cat(questions_embeds_faiss) | |
if do_normalization: | |
docs_embeds_faiss = F.normalize(docs_embeds_faiss, dim=-1) | |
questions_embeds_faiss = F.normalize(questions_embeds_faiss, dim=-1) | |
return docs_embeds_faiss.numpy(), questions_embeds_faiss.numpy() | |
def filter_ref_parts(ref_dict, filter_parts): | |
filtered_dict = {} | |
for k, refs in ref_dict.items(): | |
filtered_refs = [" ".join([x for x in ref.split() if not any([True for part in filter_parts if part in x])]) | |
for ref in refs] | |
filtered_dict[k] = filtered_refs | |
return filtered_dict | |
def get_final_metrics(pred, true, categories, top_k_values, | |
metrics_func, metrics_func_params): | |
metrics = {} | |
for top_k in top_k_values: | |
ctg_metrics = {} | |
for ctg in categories: | |
ctg_pred, ctg_true = get_exact_ctg_data(pred, true, ctg) | |
metrics_at_k = metrics_func(ctg_pred, ctg_true, top_k, **metrics_func_params) | |
for mk in metrics_at_k.keys(): | |
metrics_at_k[mk] = round(metrics_at_k[mk] * 100, 6) | |
ctg_metrics[ctg] = metrics_at_k | |
metrics[top_k] = ctg_metrics | |
return metrics | |
def get_exact_ctg_data(pred_in, true_in, ctg): | |
if ctg == "all": | |
return pred_in, true_in | |
out_pred = {} | |
out_true = {} | |
for idx, (pred, true) in zip(true_in.keys(), zip(pred_in.values(), true_in.values())): | |
ctg_refs_true = [ref for ref in true if ctg in ref] | |
ctg_refs_pred = [ref for ref in pred if ctg in ref] | |
out_true[idx] = ctg_refs_true | |
out_pred[idx] = ctg_refs_pred | |
return out_pred, out_true | |
def print_metrics(metrics, ref_categories): | |
first_ctg = metrics[list(metrics.keys())[0]] | |
metric_tags = list(first_ctg[list(first_ctg.keys())[0]].keys()) | |
metric_tags = [x.split('@')[0] for x in metric_tags] | |
print('\t', *metric_tags, sep='\t') | |
for ctg, ctg_short in ref_categories.items(): | |
for top_k, vals in metrics.items(): | |
for ctg_tag, ctg_val in vals.items(): | |
if ctg_tag == ctg: | |
ctg_vals_str = ["{:.3f}".format(x).zfill(6) for x in ctg_val.values()] | |
print(f"{ctg_short}@{top_k}", *ctg_vals_str, sep='\t\t') | |