Zekun Wu
commited on
Commit
·
d20d0a7
1
Parent(s):
99cd796
update
Browse files- 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
|
|
|
|
|
|
|
|