Update app.py
Browse files
app.py
CHANGED
@@ -3,6 +3,9 @@ from transformers import pipeline
|
|
3 |
import spaces
|
4 |
@spaces.GPU
|
5 |
# generator = pipeline('text-generation', model='distilgpt2', device=-1)
|
|
|
|
|
|
|
6 |
|
7 |
def generate_text(prompt, max_length):
|
8 |
result = generator(prompt, max_length=max_length, num_return_sequences=1)
|
|
|
3 |
import spaces
|
4 |
@spaces.GPU
|
5 |
# generator = pipeline('text-generation', model='distilgpt2', device=-1)
|
6 |
+
model_name = "rinna/japanese-gpt2-medium"
|
7 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
8 |
+
generator = pipeline('text-generation', model=model_name, tokenizer=tokenizer, device=-1)
|
9 |
|
10 |
def generate_text(prompt, max_length):
|
11 |
result = generator(prompt, max_length=max_length, num_return_sequences=1)
|