update to use my config

README.md

@@ -1,5 +1,5 @@
 ---
-title: CALM2-7B-chat
+title: chat-1
 emoji: ⚡
 colorFrom: red
 colorTo: purple

app.py

@@ -7,9 +7,9 @@ from threading import Thread
 import gradio as gr
 import spaces
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, pipeline
 
-DESCRIPTION = "# CALM2-7B-chat"
+DESCRIPTION = "# chat-1"
 
 if not torch.cuda.is_available():
     DESCRIPTION += "\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>"
@@ -20,14 +20,10 @@ MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "32768"))
 
 if torch.cuda.is_available():
     model_id = "cyberagent/calm2-7b-chat"
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-
-
-def apply_chat_template(conversation: list[dict[str, str]]) -> str:
-    prompt = "\n".join([f"{c['role']}: {c['content']}" for c in conversation])
-    return f"{prompt}\nASSISTANT: "
-
+    my_pipeline=pipeline(
+        model=model_id,
+    )
+    my_pipeline.tokenizer.chat_template = "{{bos_token}}{% for message in messages %}{% if message['role'] == 'user' %}{{ '\\n\\n### 前の投稿:\\n' + message['content'] + '' }}{% elif message['role'] == 'system' %}{{ '以下は、SNS上の投稿です。あなたはSNSの投稿生成botとして、次に続く投稿を考えなさい。説明はせず、投稿の内容のみを鉤括弧をつけずに答えよ。' }}{% elif message['role'] == 'assistant' %}{{ '\\n\\n### 次の投稿:\\n' + message['content'] + eos_token }}{% endif %}{% if loop.last and add_generation_prompt %}{{ '\\n\\n### 次の投稿:\\n' }}{% endif %}{% endfor %}"
 
 @spaces.GPU
 @torch.inference_mode()
@@ -40,38 +36,15 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.0,
 ) -> Iterator[str]:
-    conversation = []
-    for user, assistant in chat_history:
-        conversation.extend([{"role": "USER", "content": user}, {"role": "ASSISTANT", "content": assistant}])
-    conversation.append({"role": "USER", "content": message})
-
-    prompt = apply_chat_template(conversation)
-    input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    messages = [
+        {"role": "system", "content": "あなたはSNSの投稿生成botで、次に続く投稿を考えてください。"},
+        {"role": "user", "content": message},
+    ]
 
+    output = my_pipeline(
+        messages,
+    )[-1]["generated_text"][-1]["content"]
+    yield output
 
 demo = gr.ChatInterface(
     fn=generate,
@@ -116,10 +89,10 @@ demo = gr.ChatInterface(
     ],
     stop_btn=None,
     examples=[
-        ["東京の観光名所を教えて。"],
-        ["落武者って何？"],  # noqa: RUF001
-        ["暴れん坊将軍って誰のこと？"],  # noqa: RUF001
-        ["人がヘリを食べるのにかかる時間は？"],  # noqa: RUF001
+        ["サマリーを作る男の人，サマリーマン。"],
+        ["やばい場所にクリティカルな配線ができてしまったので掲示した。"],
+        ["にゃん"],
+        ["Wikipedia の情報は入っているのかもしれない"],
     ],
     description=DESCRIPTION,
     css_paths="style.css",