Update app.py

app.py

@@ -1,73 +1,48 @@
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_ID = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
+
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    load_in_8bit=True,
-    device_map="auto"
-)
+    "rinna/bilingual-gpt-neox-4b-instruction-ppo",
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
 
-import gradio as gr
+device = model.device
+device
 
-def generate_response(user_question,
-             chat_history,
+def generate(user_question,
              temperature=0.3,
-             top_p=0.85,
-             max_new_tokens=2048,
-             repetition_penalty=1.05
-            ):
-    # 挙動の指定
-    user_prompt_template = "ユーザー:あなたは日本語で質問やコメントに対して、回答してくれるアシスタントです。関西弁で回答してください"
     system_prompt_template = "システム: もちろんやで！どんどん質問してな！今日も気分ええわ！"
 
     # one-shot
+    user_sample = "ユーザー:日本でよく飲まれているお茶の種類を教えて？"
+    system_sample = "システム: 緑茶やで！緑茶って殺菌作用もあって最高よな！"
+
     user_sample = "ユーザー:日本一の高さの山は？ "
     system_sample = "システム: 富士山や！最高の眺めを拝めるで！！"
 
-    user_sample = "大阪で有名な食べ物は？ "
-    system_sample = "システム: たこ焼きやで！！外がカリカリ、中がふわふわや"
 
-    
     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):]
 
+
+    output = generate('人工知能とは何ですか？')
+output
+
 with gr.Blocks() as demo:
     chat_history = gr.Chatbot()
-    user_message = gr.Textbox(label="Question:", placeholder="人工知能とは何ですか？")
+    inputs = gr.Textbox(label="Question:", placeholder="質問を入力してください")
+    outputs = gr.Textbox(label="Answer:")
+    btn = gr.Button("Send")
     clear = gr.ClearButton([user_message, chat_history])
 
+# ボタンが押された時の動作を以下のように定義する:
+    # 「inputs内の値を入力としてモデルに渡し、その戻り値をoutputsの値として設定する」
+    btn.click(fn=generate, inputs=inputs, outputs=outputs)
+
     def response(user_message, chat_history):
-        system_message = generate_response(user_message, chat_history)
         chat_history.append((user_message, system_message))
         return "", chat_history
 
@@ -75,3 +50,5 @@ with gr.Blocks() as demo:
 
 if __name__ == "__main__":
     demo.launch()
+
+