Update app.py

app.py

@@ -1,52 +1,176 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import gradio as gr
 
+MODEL_ID = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
 model = AutoModelForCausalLM.from_pretrained(
-    "rinna/bilingual-gpt-neox-4b-instruction-ppo",
+    MODEL_ID,
+    load_in_8bit=True,
+    device_map="auto"
+)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
 
 device = model.device
-
+device
+
+user_prompt_template = "ユーザー: Hello, you are an assistant that helps me learn Japanese. I am going to ask you a question, so please answer *briefly*."
+system_prompt_template = "システム: Sure, I will answer briefly. What can I do for you?"
+
+# one-shot
+user_sample = "ユーザー: 日本で一番高い山は何ですか？"
+system_sample = "システム: 富士山です。高さは3776メートルです。"
+
+# 質問
+user_prerix = "ユーザー: "
+user_question = "人工知能とは何ですか？"
+system_prefix = "システム: "
+
+# プロンプトの整形
+prompt = user_prompt_template + "\n" + system_prompt_template + "\n"
+prompt += user_sample + "\n" + system_sample + "\n"
+prompt += user_prerix + user_question + "\n" + system_prefix
+
+inputs = tokenizer(
+    prompt,
+    add_special_tokens=False, # プロンプトに余計なトークンが付属するのを防ぐ
+    return_tensors="pt"
+)
+inputs = inputs.to(model.device)
+with torch.no_grad():
+    tokens = model.generate(
+        **inputs,
+        temperature=0.3,
+        top_p=0.85,
+        max_new_tokens=2048,
+        repetition_penalty=1.05,
+        do_sample=True,
+        pad_token_id=tokenizer.pad_token_id,
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
+tokens
+
+output = tokenizer.decode(
+    tokens[0],
+    skip_special_tokens=True # 出力に余計なトークンが付属するのを防ぐ
+)
+print(output)
+
+output[len(prompt):]
 
 def generate(user_question,
              temperature=0.3,
-    system_prompt_template = "システム: もちろんやで！どんどん質問してな！今日も気分ええわ！"
-
-    # one-shot
-
+             top_p=0.85,
+             max_new_tokens=2048,
+             repetition_penalty=1.05
+            ):
 
+    user_prompt_template = "ユーザー: Hello, you are an assistant that helps me learn Japanese. I am going to ask you a question, so please answer *briefly*."
+    system_prompt_template = "システム: Sure, I will answer briefly. What can I do for you?"
 
-    user_sample = "ユーザー:日本一の高さの山は？ "
-    system_sample = "システム: 富士山や！最高の眺めを拝めるで！！"
-
+    user_sample = "ユーザー: 日本で一番高い山は何ですか？"
+    system_sample = "システム: 富士山です。高さは3776メートルです。"
 
     user_prerix = "ユーザー: "
     system_prefix = "システム: "
 
+    prompt = user_prompt_template + "\n" + system_prompt_template + "\n"
+    prompt += user_sample + "\n" + system_sample + "\n"
+    prompt += user_prerix + user_question + "\n" + system_prefix
+
+    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
+    inputs = inputs.to(model.device)
+    with torch.no_grad():
+        tokens = model.generate(
+            **inputs,
+            temperature=temperature,
+            top_p=top_p,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            bos_token_id=tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
     output = tokenizer.decode(tokens[0], skip_special_tokens=True)
     return output[len(prompt):]
 
+output = generate('人工知能とは何ですか？')
+output
 
 
-
+import gradio as gr # 慣習としてgrと略記
 
 with gr.Blocks() as demo:
-    chat_history = gr.Chatbot()
-    inputs = gr.Textbox(label="Question:", placeholder="質問を入力してください")
+    inputs = gr.Textbox(label="Question:", placeholder="人工知能とは何ですか？")
     outputs = gr.Textbox(label="Answer:")
     btn = gr.Button("Send")
-    clear = gr.ClearButton([inputs, chat_history])
 
     # ボタンが押された時の動作を以下のように定義する:
-
+    # 「inputs内の値を入力としてモデルに渡し、その戻り値をoutputsの値として設定する」
     btn.click(fn=generate, inputs=inputs, outputs=outputs)
 
+if __name__ == "__main__":
+    demo.launch()
+    
+def generate_response(user_question,
+             chat_history,
+             temperature=0.3,
+             top_p=0.85,
+             max_new_tokens=2048,
+             repetition_penalty=1.05
+            ):
+
+    user_prompt_template = "ユーザー: Hello, you are an assistant that helps me learn Japanese. I am going to ask you a question, so please answer *briefly*."
+    system_prompt_template = "システム: Sure, I will answer briefly. What can I do for you?"
+
+    user_sample = "ユーザー: 日本で一番高い山は何ですか？"
+    system_sample = "システム: 富士山です。高さは3776メートルです。"
+
+    user_prerix = "ユーザー: "
+    system_prefix = "システム: "
+
+    prompt = user_prompt_template + "\n" + system_prompt_template + "\n"
+
+    if len(chat_history) < 1:
+        prompt += user_sample + "\n" + system_sample + "\n"
+    else:
+        u = chat_history[-1][0]
+        s = chat_history[-1][1]
+        prompt += user_prerix + u + "\n" + system_prefix + s + "\n"
+
+    prompt += user_prerix + user_question + "\n" + system_prefix
+
+    inputs = tokenizer(prompt, add_special_tokens=False, return_tensors="pt")
+    inputs = inputs.to(model.device)
+    with torch.no_grad():
+        tokens = model.generate(
+            **inputs,
+            temperature=temperature,
+            top_p=top_p,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            pad_token_id=tokenizer.pad_token_id,
+            bos_token_id=tokenizer.bos_token_id,
+            eos_token_id=tokenizer.eos_token_id
+        )
+    output = tokenizer.decode(tokens[0], skip_special_tokens=True)
+    return output[len(prompt):]
+
+
+with gr.Blocks() as demo:
+    chat_history = gr.Chatbot()
+    user_message = gr.Textbox(label="Question:", placeholder="人工知能とは何ですか？")
+    clear = gr.ClearButton([user_message, chat_history])
+
     def response(user_message, chat_history):
+        system_message = generate_response(user_message, chat_history)
         chat_history.append((user_message, system_message))
         return "", chat_history
 
-    inputs.submit(response, inputs=[inputs, chat_history], outputs=[inputs, chat_history])
+    user_message.submit(response, inputs=[user_message, chat_history], outputs=[user_message, chat_history])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()
+