Views
No views yet
1import torch
2from transformers import T5Tokenizer, T5ForConditionalGeneration
3
4trained_model_path = 'ZhangCheng/T5-Base-Fine-Tuned-for-Question-Generation'
5trained_tokenizer_path = 'ZhangCheng/T5-Base-Fine-Tuned-for-Question-Generation'
6
7class QuestionGeneration:
8
9 def __init__(self, model_dir=None):
10 self.model = T5ForConditionalGeneration.from_pretrained(trained_model_path)
11 self.tokenizer = T5Tokenizer.from_pretrained(trained_tokenizer_path)
12 self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
13 self.model = self.model.to(self.device)
14 self.model.eval()
15
16 def generate(self, answer: str, context: str):
17 input_text = '<answer> %s <context> %s ' % (answer, context)
18 encoding = self.tokenizer.encode_plus(
19 input_text,
20 return_tensors='pt'
21 )
22 input_ids = encoding['input_ids']
23 attention_mask = encoding['attention_mask']
24 outputs = self.model.generate(
25 input_ids=input_ids,
26 attention_mask=attention_mask
27 )
28 question = self.tokenizer.decode(
29 outputs[0],
30 skip_special_tokens=True,
31 clean_up_tokenization_spaces=True
32 )
33 return {'question': question, 'answer': answer, 'context': context}
34
35if __name__ == "__main__":
36 context = 'ZhangCheng fine-tuned T5 on SQuAD dataset for question generation.'
37 answer = 'ZhangCheng'
38 QG = QuestionGeneration()
39 qa = QG.generate(answer, context)
40 print(qa['question'])
41 # Output:
42 # Who fine-tuned T5 on SQuAD dataset for question generation?