change model

app.py

@@ -26,56 +26,26 @@ import gradio as gr
 
 # gr.ChatInterface(predict).launch()
 
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseTokenizer
-# トークナイザーとモデルの準備
-#tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt2-medium")
-#model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt2-medium")
-model = GPTNeoXJapaneseForCausalLM.from_pretrained("abeja/gpt-neox-japanese-2.7b")
-tokenizer = GPTNeoXJapaneseTokenizer.from_pretrained("abeja/gpt-neox-japanese-2.7b")
-#model = AutoModelForCausalLM.from_pretrained("cl-tohoku/bert-base-japanese")
-#tokenizer = AutoTokenizer.from_pretrained("cl-tohoku/bert-base-japanese")
+from transformers import pipeline
+from transformers import Conversation
 
-# user関数 (history:セッション状態)
-def user(message, history):
-    return "", history + [[message, None]]
-
-# bot関数 (history:セッション状態)
-def bot(history):
-    user_message = history[-1][0]
-    print(user_message)
-    print(tokenizer.eos_token)
-    new_user_input_ids = tokenizer.encode(
-        user_message + tokenizer.eos_token, return_tensors="pt"
-    )
+import gradio as gr
 
-    # レスポンスの生成
-    bot_input_ids = torch.cat([torch.LongTensor([]), new_user_input_ids], dim=-1)
-    response = model.generate(
-        bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id
-    ).tolist()
+chatbot = pipeline(model="facebook/blenderbot-400M-distill")
+message_list = []
+response_list = []
 
-    # トークンをテキストに変換し、応答を行に分割
-    response = tokenizer.decode(response[0]).split("<|endoftext|>")
-    response = [
-        (response[i], response[i + 1]) for i in range(0, len(response) - 1, 2)
-    ]
-    history[-1] = response[0]
-    return history
+def mini_chatbot(message, history):
+    conversation = Conversation(text=message,
+                                past_user_inputs=message_list,
+                                generated_responses=response_list)
+    conversation = chatbot(conversation)
 
-# Blocksの作成
-with gr.Blocks() as demo:
-    # UI
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox()
-    clear = gr.Button("Clear")
+    return conversation.generated_responses[-1]
 
-    # イベントリスナー
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
+demo_chatbot = gr.ChatInterface(mini_chatbot,
+                                title="My Chatbot",
+                                description="Enter text to start chatting.")
+                                
 
-# 起動
-demo.launch()
+demo_chatbot.launch()