Views
No views yet
load_finetuned_model to load the query and sentence encoder, and encode_batch() to encode a sentence with the model.1
2from transformers import AutoTokenizer, AutoModel
3import torch
4
5def load_finetuned_model():
6
7
8 sentence_encoder = AutoModel.from_pretrained("biu-nlp/abstract-sim-sentence-pubmed", revision="71f4539120e29024adc618173a1ed5fd230ac249")
9 query_encoder = AutoModel.from_pretrained("biu-nlp/abstract-sim-query-pubmed", revision="8d34676d80a39bcbc5a1d2eec13e6f8078496215")
10 tokenizer = AutoTokenizer.from_pretrained("biu-nlp/abstract-sim-sentence-pubmed")
11 return tokenizer, query_encoder, sentence_encoder
12
13
14def encode_batch(model, tokenizer, sentences, device):
15 input_ids = tokenizer(sentences, padding=True, max_length=128, truncation=True, return_tensors="pt",
16 add_special_tokens=True).to(device)
17 features = model(**input_ids)[0]
18
19 features = torch.sum(features[:,:,:] * input_ids["attention_mask"][:,:].unsqueeze(-1), dim=1) / torch.clamp(torch.sum(input_ids["attention_mask"][:,:], dim=1, keepdims=True), min=1e-9)
20
21 return features
22