Views
No views yet
1from transformers import AutoTokenizer, AutoModel
2
3# load model and tokenizer
4tokenizer = AutoTokenizer.from_pretrained('malteos/PubMedNCL')
5model = AutoModel.from_pretrained('malteos/PubMedNCL')
6
7papers = [{'title': 'BERT', 'abstract': 'We introduce a new language representation model called BERT'},
8 {'title': 'Attention is all you need', 'abstract': ' The dominant sequence transduction models are based on complex recurrent or convolutional neural networks'}]
9
10# concatenate title and abstract with [SEP] token
11title_abs = [d['title'] + tokenizer.sep_token + (d.get('abstract') or '') for d in papers]
12
13# preprocess the input
14inputs = tokenizer(title_abs, padding=True, truncation=True, return_tensors="pt", max_length=512)
15
16# inference
17result = model(**inputs)
18
19# take the first token ([CLS] token) in the batch as the embedding
20embeddings = result.last_hidden_state[:, 0, :]