tamatwi commited on
Commit
8e311d6
1 Parent(s): 33a37b5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -1,12 +1,12 @@
1
  import gradio as gr
2
  from transformers import pipeline, AutoTokenizer
3
- # @spaces.GPU
4
  # 日本語モデルを指定
5
  model_name = "rinna/japanese-gpt2-medium"
6
 
7
  # トークナイザーとパイプラインの設定
8
  tokenizer = AutoTokenizer.from_pretrained(model_name)
9
- generator = pipeline('text-generation', model=model_name, tokenizer=tokenizer, device=-1)
10
 
11
  def generate_text(prompt, max_length):
12
  result = generator(prompt, max_length=max_length, num_return_sequences=1)
 
1
  import gradio as gr
2
  from transformers import pipeline, AutoTokenizer
3
+
4
  # 日本語モデルを指定
5
  model_name = "rinna/japanese-gpt2-medium"
6
 
7
  # トークナイザーとパイプラインの設定
8
  tokenizer = AutoTokenizer.from_pretrained(model_name)
9
+ generator = pipeline('text-generation', model=model_name, tokenizer=tokenizer, device=0) # device=0はGPUを使用する設定
10
 
11
  def generate_text(prompt, max_length):
12
  result = generator(prompt, max_length=max_length, num_return_sequences=1)