pytorch
Browse files
app.py
CHANGED
@@ -1,7 +1,7 @@
|
|
1 |
import gradio as gr
|
2 |
|
3 |
import transformers
|
4 |
-
from transformers import AutoTokenizer,
|
5 |
import tensorflow as tf
|
6 |
|
7 |
SAVED_CHECKPOINT = 'mikegarts/distilgpt2-erichmariaremarque'
|
@@ -9,7 +9,7 @@ MIN_WORDS = 80
|
|
9 |
|
10 |
|
11 |
def get_model():
|
12 |
-
model =
|
13 |
tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
|
14 |
return model, tokenizer
|
15 |
|
@@ -18,14 +18,15 @@ def generate(prompt):
|
|
18 |
model, tokenizer = get_model()
|
19 |
|
20 |
input_context = prompt
|
21 |
-
input_ids = tokenizer.encode(input_context, return_tensors="
|
22 |
|
23 |
outputs = model.generate(
|
24 |
-
input_ids=input_ids,
|
25 |
-
max_length=
|
26 |
-
temperature=0.7,
|
27 |
-
num_return_sequences=
|
28 |
-
do_sample=True
|
|
|
29 |
)
|
30 |
|
31 |
return tokenizer.decode(outputs[0], skip_special_tokens=True).rsplit('.', 1)[0] + '.'
|
|
|
1 |
import gradio as gr
|
2 |
|
3 |
import transformers
|
4 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
5 |
import tensorflow as tf
|
6 |
|
7 |
SAVED_CHECKPOINT = 'mikegarts/distilgpt2-erichmariaremarque'
|
|
|
9 |
|
10 |
|
11 |
def get_model():
|
12 |
+
model = AutoModelForCausalLM.from_pretrained(SAVED_CHECKPOINT)
|
13 |
tokenizer = AutoTokenizer.from_pretrained(SAVED_CHECKPOINT)
|
14 |
return model, tokenizer
|
15 |
|
|
|
18 |
model, tokenizer = get_model()
|
19 |
|
20 |
input_context = prompt
|
21 |
+
input_ids = tokenizer.encode(input_context, return_tensors="pt").to('cuda')
|
22 |
|
23 |
outputs = model.generate(
|
24 |
+
input_ids=input_ids,
|
25 |
+
max_length=100,
|
26 |
+
temperature=0.7,
|
27 |
+
num_return_sequences=3,
|
28 |
+
do_sample=True,
|
29 |
+
# forced_eos_token_id=tokenizer.encode('.')[0]
|
30 |
)
|
31 |
|
32 |
return tokenizer.decode(outputs[0], skip_special_tokens=True).rsplit('.', 1)[0] + '.'
|