Update app.py

app.py

@@ -1,41 +1,30 @@
 import os
 os.system("pip install transformers")
-from transformers import AutoModel, AutoTokenizer, trainer_utils
 import gradio as gr
+from transformers import AutoModel, AutoTokenizer, trainer_utils
+import torch
 
-def generate(prefix_text):
-    device = "cuda"
-
-    model = AutoModel.from_pretrained("Tanrei/GPTSAN-japanese").to(device)
-
-    tokenizer = AutoTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
-
-    x_token = tokenizer("", prefix_text=prefix_text, return_tensors="pt")
+# Load the model and tokenizer
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModel.from_pretrained("Tanrei/GPTSAN-japanese").to(device)
+tokenizer = AutoTokenizer.from_pretrained("Tanrei/GPTSAN-japanese")
 
+# Define the text generation function
+def generate_text(input_text):
+    x_token = tokenizer(input_text, return_tensors="pt")
     trainer_utils.set_seed(30)
-
     input_ids = x_token.input_ids.to(device)
-
-    token_type_ids = x_token.token_type_ids.to(device)
-
-    gen_token = model.generate(input_ids, token_type_ids=token_type_ids, max_new_tokens=50)
-
+    gen_token = model.generate(input_ids, max_new_tokens=50)
     return tokenizer.decode(gen_token[0])
 
-
-def main():
-    # テキストエリアを追加
-    text_input = gr.Textbox()
-
-    # 送信ボタンを追加
-    submit_button = gr.Button("送信")
-
-    # テキストエリアと送信ボタンを連携
-    submit_button.on_click(lambda: gr.show(generate(text_input.value), title="テキスト生成"))
-
-    # アプリを起動
-    gr.Interface(main).launch()
-
-
-if __name__ == "__main__":
-    main()
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs="text",
+    outputs="text",
+    title="Japanese Text Generation",
+    description="Enter a prompt in Japanese to generate text."
+)
+
+# Launch the interface
+iface.launch()