Spaces:
Running
Running
| import pandas as pd | |
| import torch | |
| from transformers import AutoTokenizer, AutoModel, set_seed | |
| from tqdm import tqdm | |
| from utils import clean_text | |
| from model import MimicTransformer | |
| set_seed(42) | |
| def read_model(model, path): | |
| model.load_state_dict(torch.load(path, map_location=torch.device('cuda')), strict=False) | |
| return model | |
| model_path = 'checkpoint_0_9113.bin' | |
| mimic = MimicTransformer(cutoff=512) | |
| mimic = read_model(model=mimic, path=model_path) | |
| mimic.eval() | |
| mimic.cuda() | |
| tokenizer = mimic.tokenizer | |
| summaries = pd.read_csv('all_summaries_backup.csv')['SUMMARIES'] | |
| def mean_pooling(model_output, attention_mask): | |
| token_embeddings = model_output[0] #First element of model_output contains all token embeddings | |
| input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float() | |
| return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9) | |
| def get_model_outputs(text): | |
| inputs = tokenizer(text, return_tensors='pt', padding='max_length', max_length=512, truncation=True).to('cuda') | |
| outputs = mimic(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask, drg_labels=None) | |
| # pooled = mean_pooling(outputs[0][0], inputs['attention_mask']) | |
| pooled = outputs[0][0] | |
| normalized = pooled/pooled.norm(dim=1)[:,None] | |
| return normalized | |
| return_tensors = torch.zeros(size=(10000, 738)) | |
| non_defunct_summaries = [] | |
| for i, summary in tqdm(enumerate(summaries[:50000])): | |
| cleaned = clean_text(summary) | |
| if len(non_defunct_summaries) == 10000: | |
| break | |
| if len(cleaned) > 100: | |
| non_defunct_summaries.append(cleaned) | |
| for i, summary in tqdm(enumerate(non_defunct_summaries)): | |
| res = get_model_outputs(text=summary) | |
| return_tensors[i, :] = res.detach().cpu() | |
| # sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask']) | |
| # sentence_embeddings = sentence_embeddings/sentence_embeddings.norm(dim=1)[:,None] | |
| pd.DataFrame(data={'SUMMARIES':non_defunct_summaries}).to_csv('all_summaries.csv', index=False) | |
| torch.save(return_tensors, f='discharge_embeddings.pt') | |