Update app.py

app.py

@@ -3,14 +3,23 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 import torch
 from threading import Thread
 
-MODEL_ID = "HODACHI/EZO-Common-9B-gemma-2-it"
+MODEL_ID = "HODACHI/Llama-3.1-8B-EZO-1.1-it"
 DTYPE = torch.bfloat16
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    device_map="cuda",
-    torch_dtype=DTYPE,
+    torch_dtype=torch.bfloat16,  # bfloat16形式で計算を行い、精度と速度のバランスを取る
+    device_map="auto",  # 利用可能なデバイスに自動的にモデルを配置
+    low_cpu_mem_usage=True,  # CPU消費メモリを抑える
+)
+
+pipeline = transformers.pipeline(
+    "text-generation",  # タスクを指定（ここではテキスト生成）
+    model=model,  # 使用するモデル
+    tokenizer=tokenizer,  # 使用するトークナイザー
+    device_map="auto",  # デバイスの自動割り当て
 )
 
 def respond(
@@ -21,32 +30,42 @@ def respond(
     top_p,
 ):
     chat = []
+    chat.append({"role": "system", "content": "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、原則日本語で回答してください。"})
     for user, assistant in history:
         chat.append({"role": "user", "content": user})
         chat.append({"role": "assistant", "content": assistant})
     chat.append({"role": "user", "content": message})
 
-    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    #inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt").to(model.device)
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
     
-    generation_kwargs = dict(
-        input_ids=inputs,
-        max_new_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-        do_sample=True,
-        streamer=streamer,
-    )
+    #generation_kwargs = dict(
+    #    input_ids=inputs,
+    #    max_new_tokens=max_tokens,
+    #    temperature=temperature,
+    #    top_p=top_p,
+    #    do_sample=True,
+    #    streamer=streamer,
+    #)
     
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
-    response = ""
-    for new_text in streamer:
-        response += new_text
-        yield response
+    #response = ""
+    #for new_text in streamer:
+    #    response += new_text
+    #    yield response
+    outputs = pipeline(
+        prompt,
+        max_new_tokens=40,  # 生成する最大トークン数
+        do_sample=True,  # サンプリングを有効にして多様な出力を得る
+        temperature=0.7,  # 生成の多様性を調整（高いほど多様、低いほど決定的）
+        top_p=0.95,  # 累積確率に基づくサンプリングの閾値
+    )
+
+    response = outputs[0]["generated_text"]
 
 demo = gr.ChatInterface(
     respond,