Update app.py

app.py

@@ -5,7 +5,8 @@ import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import os
 from threading import Thread
-
+import random
+from datasets import load_dataset
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL_ID = "TeamDelta/mistral-yuki-7B"
@@ -42,14 +43,19 @@ h3 {
 }
 """
 
-
+# モデルとトークナイザーの読み込み
 model = AutoModelForCausalLM.from_pretrained(
-          MODEL_ID,
-          torch_dtype=torch.float16,
-          device_map="auto",
-        )
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
+# データセットをロードして10個の例を取得
+dataset = load_dataset("elyza/ELYZA-tasks-100")
+examples = random.sample(dataset['train'], 10)
+example_inputs = [example['input'] for example in examples]
+
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
     print(f'message is - {message}')
@@ -59,8 +65,6 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         conversation.extend([{"role": "user", "content": prompt}, {"role": "assistant", "content": answer}])
     conversation.append({"role": "user", "content": message})
 
-    #print(f"Conversation is -\n{conversation}")
-    
     input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(input_ids, return_tensors="pt").to(0)
     
@@ -75,7 +79,7 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         max_new_tokens=max_new_tokens, 
         do_sample=True, 
         temperature=temperature,
-        eos_token_id = [128001, 128009],
+        eos_token_id=[128001, 128009],
     )
     
     thread = Thread(target=model.generate, kwargs=generate_kwargs)
@@ -86,8 +90,6 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
         buffer += new_text
         yield buffer
 
-
-
 chatbot = gr.Chatbot(height=500)
 
 with gr.Blocks(css=CSS) as demo:
@@ -144,16 +146,9 @@ with gr.Blocks(css=CSS) as demo:
                 render=False,
             ),
         ],
-        examples=[
-            ["超能力を持つ主人公のSF物語のシナリオを考えてください。伏線の設定、テーマやログラインを理論的に使用してください"],
-            ["子供の夏休みの自由研究のための、5つのアイデアと、その手法を簡潔に教えてください。"],
-            ["パズルゲームのスクリプト作成のためにアドバイスお願いします"],
-            ["マークダウン記法にて、ブロック崩しのゲーム作成の教科書作成してください"],
-        ],
+        examples=example_inputs,
         cache_examples=False,
     )
 
-
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()