Zekun Wu commited on
Commit
d20d0a7
1 Parent(s): 99cd796
Files changed (1) hide show
  1. utils/model.py +0 -9
utils/model.py CHANGED
@@ -6,14 +6,5 @@ class gpt2:
6
  self.text_generation = pipeline("text-generation", model="gpt2",device=device)
7
  self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
8
 
9
- def generate_text(self,**kwargs):
10
- results = self.text_generation(**kwargs)
11
-
12
- return [item['generated_text'] for item in results[0]]
13
-
14
  def get_tokenizer(self):
15
  return self.tokenizer
16
-
17
- if __name__ == '__main__':
18
- gpt2 = gpt2()
19
- print(gpt2.generate_text(["Hello, how are you?","I am fine, thank you."]))
 
6
  self.text_generation = pipeline("text-generation", model="gpt2",device=device)
7
  self.tokenizer = AutoTokenizer.from_pretrained("gpt2")
8
 
 
 
 
 
 
9
  def get_tokenizer(self):
10
  return self.tokenizer