Views
No views yet
1from datasets import load_dataset
2from transformers import AutoTokenizer, AutoModel
3from torch.utils.data import DataLoader
4import torch
5import pandas as pd
6
7# choose GPU when available
8device = 'cuda' if torch.cuda.is_available() else 'cpu'
9
10tokenizer = AutoTokenizer.from_pretrained("bertin-project/bertin-roberta-base-spanish",model_max_length=512)
11
12# build custom model with classification layer on top and a dropout layer before
13class RobertaClass(torch.nn.Module):
14
15 def __init__(self):
16 super(RobertaClass, self).__init__()
17 self.l1 = AutoModel.from_pretrained("bertin-project/bertin-roberta-base-spanish",return_dict=False)
18 self.l2 = torch.nn.Dropout(0.3)
19 self.l3 = torch.nn.Linear(768, 11)
20
21 def forward(self, input_ids, attention_mask):
22 _, output_1= self.l1(input_ids=input_ids, attention_mask=attention_mask)
23 output_2 = self.l2(output_1)
24 output = self.l3(output_2)
25
26 return output
27
28model_name="bertin-roberta-base-spanish_semeval18_emodetection/pytorch_model.bin"
29
30model=RobertaClass()
31
32model.load_state_dict(torch.load(model_name,map_location=torch.device(device)))
33
34model.eval()
35
36# run on more than 1 GPU
37model = torch.nn.DataParallel(model)
38
39model.to(device)
40
41twnames=['anger','anticipation','disgust','fear','joy','love','optimism','pessimism','sadness','surprise','trust']
42
43# load from hugging face dataset hub
44testset_raw = load_dataset('sem_eval_2018_task_1','subtask5.spanish',split='test')
45
46# remove old columns
47testset=testset_raw.remove_columns(twnames+["ID"])
48
49# tokenize
50testset_tokenized = testset.map(lambda e: tokenizer(e['Tweet'], truncation=True, padding='max_length'), batched=True)
51
52testset_tokenized=testset_tokenized.remove_columns("Tweet")
53
54testset_tokenized.set_format(type='torch', columns=['input_ids', 'attention_mask'])
55
56
57outfile="predicted_2018-E-c-Es-test-gold.txt"
58
59MAX_LEN = 512
60VALID_BATCH_SIZE = 8
61# set batch size according to available RAM
62# VALID_BATCH_SIZE = 1000
63
64# set num_workers for parallel processing
65inference_params = {'batch_size': VALID_BATCH_SIZE,
66 'shuffle': False,
67 # 'num_workers': 1
68 }
69
70inference_loader = DataLoader(testset_tokenized, **inference_params)
71
72
73open(outfile,"w").close()
74with torch.no_grad():
75 # change lines for progress manager
76 # for _, data in tqdm(enumerate(inference_loader, 0),total=len(inference_loader)):
77 for _, data in enumerate(inference_loader, 0):
78 outputs = model(input_ids=data['input_ids'],attention_mask=data['attention_mask'])
79 fin_outputs=torch.sigmoid(outputs).cpu().detach().numpy().tolist()
80 pd.DataFrame(fin_outputs).to_csv(outfile,index=False,header=False,sep="\t",mode='a')
81
82
83# # dataset from file (one text per line)
84# from datasets import Dataset
85
86# with open(linesoftextfile,"rb") as textfile:
87# textdict={"text":[x.decode().rstrip("\n") for x in textfile.readlines()]}
88
89# inference_dataset=Dataset.from_dict(textdict)
90# del(textdict)