Update app.py

app.py

@@ -3,8 +3,8 @@ from transformers import pipeline, AutoTokenizer
 import torch
 import gradio as gr
 
-# PersonaHubデータセットの読み込み
-dataset = load_dataset("proj-persona/PersonaHub", "math")
+# PersonaHubデータセットの読み込み（'math'設定を指定）
+dataset = load_dataset("proj-persona/PersonaHub", "persona")
 
 # 日本語モデルを指定
 model_name = "rinna/japanese-gpt2-medium"
@@ -21,9 +21,12 @@ def generate_text(prompt, max_length, persona_id):
     # ペルソナのプロンプトを作成
     persona_prompt = f"あなたは{persona['profession']}です。次の問題を作成してください："
     
-    # ペルソナのプロンプトを追加
+    # ペルソナのプロンプトとユーザー入力を結合
     full_prompt = f"{persona_prompt} {prompt}"
+    
+    # テキストを生成
     result = generator(full_prompt, max_length=max_length, num_return_sequences=1)
+    
     return result[0]['generated_text']
 
 iface = gr.Interface(