Update app.py

app.py

@@ -38,7 +38,7 @@ h3 {
 .chatbox .messages .message.user {
     background-color: #e1f5fe;
 }
-.chatbox .messages .message.bot {
+chatbox .messages .message.bot {
     background-color: #eeeeee;
 }
 """
@@ -51,9 +51,15 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# データセットをロードして10個の例を取得
+# データセットをロードしてスプリットを確認
 dataset = load_dataset("elyza/ELYZA-tasks-100")
-examples = random.sample(dataset['train'], 10)
+print(dataset)
+
+# 使用するスプリット名を確認
+split_name = "train" if "train" in dataset else "test"  # デフォルトをtrainにし、なければtestにフォールバック
+
+# 適切なスプリットから10個の例を取得
+examples = random.sample(dataset[split_name], 10)
 example_inputs = [example['input'] for example in examples]
 
 @spaces.GPU