Update app.py

app.py

@@ -1,86 +1,63 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline
 import torch
 from threading import Thread
-import transformers
 
 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,
-    torch_dtype=torch.bfloat16,  # bfloat16形式で計算を行い、精度と速度のバランスを取る
-    device_map="auto",  # 利用可能なデバイスに自動的にモデルを配置
-    low_cpu_mem_usage=True,  # CPU消費メモリを抑える
+    torch_dtype=DTYPE,
+    device_map="auto",
+    low_cpu_mem_usage=True,
 )
 
-pipeline = transformers.pipeline(
-    "text-generation",  # タスクを指定（ここではテキスト生成）
-    model=model,  # 使用するモデル
-    tokenizer=tokenizer,  # 使用するトークナイザー
-    device_map="auto",  # デバイスの自動割り当て
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto",
 )
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    max_tokens,
-    temperature,
-    top_p,
-):
+def generate_text(prompt, max_new_tokens, temperature, top_p):
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(
+        max_new_tokens=max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        streamer=streamer,
+    )
+    
+    thread = Thread(target=pipe, kwargs=dict(text_inputs=prompt, **generation_kwargs))
+    thread.start()
+    
+    return streamer
+
+def respond(message, history, max_tokens, temperature, 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)
-
-    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,
-    #)
+    streamer = generate_text(prompt, max_tokens, temperature, top_p)
     
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
-
-    #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"]
-    return response
+    response = ""
+    for new_text in streamer:
+        response += new_text
+        yield response
 
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
         gr.Slider(minimum=1, maximum=2048, value=150, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
 )