Views
No views yet
1class T5ForSQG:
2 def __init__(self, model_path):
3 self.model = T5ForConditionalGeneration.from_pretrained(model_path)
4 self.tokenizer = T5Tokenizer.from_pretrained(model_path)
5
6 def make_queries(self, topic, n=1, device='cpu', batch_size=16):
7 ds = YourDataSetClass(pd.DataFrame({'topic': ['make queries: '+topic]*n, 'queries': [[]*n]}, index=range(n)), self.tokenizer, 64, 64, 'topic', 'queries')
8
9 loader_params = {'batch_size': n if n < batch_size else batch_size, 'shuffle': False, 'num_workers': 0}
10
11 loader = DataLoader(ds, **loader_params)
12
13 self.model.eval()
14
15 predictions = []
16 with torch.no_grad():
17 for _, data in enumerate(loader, 0):
18 y = data['target_ids'].to(device, dtype = torch.long)
19 ids = data['source_ids'].to(device, dtype = torch.long)
20 mask = data['source_mask'].to(device, dtype = torch.long)
21
22 generated_ids = self.model.generate(
23 input_ids = ids,
24 attention_mask = mask,
25 max_length=64,
26 num_beams=1,
27 repetition_penalty=2.5,
28 length_penalty=1.0,
29 do_sample = True,
30 temperature = 1.5,
31 top_k = 10,
32 top_p = 0.95
33 )
34
35 preds = list(set([self.tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True) for g in generated_ids]))
36 predictions.extend(preds)
37
38 return list(set(predictions))