Upload 3 files

ChatSystem.py

@@ -0,0 +1,53 @@
+import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+class RinnaTalk():
+    def __init__(self, tokenizer=None, model=None):
+        self.prompt = ''
+        # 事前にmodelとtokenizerを呼んでおく
+        self.tokenizer = AutoTokenizer.from_pretrained("Hoshiyume/FixedStar-BETA", use_fast=False) if tokenizer is None else tokenizer
+        self.model = AutoModelForCausalLM.from_pretrained("Hoshiyume/FixedStar-BETA", torch_dtype=torch.float16) if model is None else model
+
+    def chat(self, message: str, chat_history: list, max_token_length: int = 128, min_token_length: int = 10, temperature: float = 0.8):
+        # チャット履歴をクリアした際にpromptもクリアさせるため
+        if len(chat_history) == 0:
+            self.prompt = ''
+
+        self.prompt += f'ユーザー: {message}\nシステム: '
+        token_ids = self.tokenizer.encode(self.prompt, add_special_tokens=False, return_tensors="pt")
+        with torch.no_grad():
+            output_ids = self.model.generate(
+                token_ids,
+                max_new_tokens=max_token_length,
+                min_new_tokens=min_token_length,
+                top_p=top_p,
+                top_k=top_k,
+                do_sample=do_sample,
+                temperature=temperature,
+                num_beam=num_beam,
+                pad_token_id=self.tokenizer.pad_token_id,
+                bos_token_id=self.tokenizer.bos_token_id,
+                eos_token_id=self.tokenizer.eos_token_id
+            )
+        output = self.tokenizer.decode(output_ids.tolist()[0])
+        latest_reply = output.split('<NL>')[-1].rstrip('</s>')
+        chat_history.append([message, latest_reply])
+        self.prompt += f'{latest_reply}\n'
+
+        return "", chat_history
+rinna = RinnaTalk()
+
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot( label="FixedStar-DebugChat", show_copy_button=True, show_share_button=True, avatar_images=["user-icon.png", "FixedStar-icon.png"] )
+    max_token_length = gr.Slider( value=512, minimum=10, maximum=512, label='max_token_length' )
+    min_token_length = gr.Slider( value=1, minimum=1, maximum=512, label='min_token_length' )
+    top_p = gr.Slider( value=0.75, minimum=0, maximum=1, label='top_p' )
+    top_k = gr.Slider( value=40, minimum=1, maximum=1000, label='top_k' )
+    temperature = gr.Slider( value=0.9, minimum=0, maximum=1, scale=0.01, value=0.8, label='temperature' )
+    do_sample = gr.Checkbox( value=True, label='do_sample' )
+    num_beam = gr.Slider( value=0, minimum=0, maximum=100, label='num_beam' )
+    msg = gr.Textbox()
+    clear = gr.ClearButton([msg, chatbot])
+    msg.submit(rinna.chat, [msg, chatbot, max_token_length, min_token_length, temperature], [msg, chatbot])
+demo.launch()