Upload app2.py

app2.py

@@ -0,0 +1,209 @@
+from typing import Any, List, Union, Mapping, Optional, Iterable
+
+import ctranslate2
+from ctranslate2 import GenerationStepResult, Generator
+import transformers
+from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast, AutoTokenizer
+
+from langchain.llms.base import LLM
+from langchain.callbacks.manager import CallbackManagerForLLMRun
+
+generator=Generator("E:\datafile\AI\FixedStar-model\FixedStar-BETA-7b-ct2")
+tokenizer=AutoTokenizer.from_pretrained("E:\datafile\AI\FixedStar-model\FixedStar-BETA-7b")
+TITLE = "Chat room!"
+
+class CTranslate2StreamLLM(LLM):
+
+    generator: Generator
+    tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]
+
+    max_length: int = 128
+    repetition_penalty: float = 1.1
+    temperature: float = 0.6
+    topk: int = 1
+
+    @property
+    def _llm_type(self) -> str:
+        return "CTranslate2"
+
+    def _generate_tokens(
+        self,
+        prompt: str,
+    ) -> Iterable[GenerationStepResult]:
+        
+        # 推論の実行
+        tokens = self.tokenizer.convert_ids_to_tokens(
+            self.tokenizer.encode(prompt, add_special_tokens=False)
+        )
+
+        step_results = self.generator.generate_tokens(
+            tokens,
+            max_length=self.max_length,
+            sampling_topk=self.topk,
+            sampling_temperature=self.temperature,
+            repetition_penalty=self.repetition_penalty,
+            return_log_prob=True,
+            end_token=[26168, 27, 208, 14719, 9078, 18482, 27, 208],
+        )
+
+        return step_results
+
+    def _decode_with_buffer(
+        self, step_result: GenerationStepResult, token_buffer: list
+    ) -> Union[str, None]:
+        token_buffer.append(step_result.token_id)
+        word = self.tokenizer.decode(token_buffer)
+
+        # 全て変換不能文字の場合、終了
+        if all(c == "�" for c in word):
+            return None
+
+        # step_resultのtokenが▁から始まる場合、スペースを付与する
+        if step_result.token.startswith("▁"):
+            word = " " + word
+
+        # 正常な文字が生成できた場合、バッファをクリア
+        token_buffer.clear()
+
+        return word
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+    ) -> str:
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+
+        step_results = self._generate_tokens(prompt)
+
+        output_ids = []
+        token_buffer = []
+
+        for step_result in step_results:
+
+            output_ids.append(step_result.token_id)
+
+            if run_manager:
+                if word := self._decode_with_buffer(step_result, token_buffer):
+
+                    run_manager.on_llm_new_token(
+                        word,
+                        verbose=self.verbose,
+                        logprobs=step_result.log_prob if step_result.log_prob else None,
+                    )
+
+        if output_ids:
+            text = self.tokenizer.decode(output_ids)
+            return text
+
+        return ""
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+    ) -> str:
+        if stop is not None:
+            raise ValueError("stop kwargs are not permitted.")
+
+        step_results = self._generate_tokens(prompt)
+
+        output_ids = []
+        token_buffer = []
+
+        for step_result in step_results:
+
+            output_ids.append(step_result.token_id)
+
+            if run_manager:
+                if word := self._decode_with_buffer(step_result, token_buffer):
+
+                    await run_manager.on_llm_new_token(
+                        word,
+                        verbose=self.verbose,
+                        logprobs=step_result.log_prob if step_result.log_prob else None,
+                    )
+
+        if output_ids:
+            text = self.tokenizer.decode(output_ids)
+            return text
+
+        return ""
+
+    @property
+    def _identifying_params(self) -> Mapping[str, Any]:
+        """Get the identifying parameters."""
+        return {
+            "generator": self.generator,
+            "tokenizer": self.tokenizer,
+            "max_length": self.max_length,
+            "repetition_penalty": self.repetition_penalty,
+            "temperature": self.temperature,
+            "topk": self.topk,
+        }
+    
+from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+llm = CTranslate2StreamLLM(
+    generator=generator,
+    tokenizer=tokenizer,
+    callbacks=[StreamingStdOutCallbackHandler()])
+
+
+# ウェブUIの起動
+import os
+import itertools
+import gradio as gr
+
+def make_prompt(message, chat_history, max_context_size: int = 10):
+    contexts = chat_history + [[message, ""]]
+    contexts = list(itertools.chain.from_iterable(contexts))
+    if max_context_size > 0:
+        context_size = max_context_size - 1
+    else:
+        context_size = 100000
+    contexts = contexts[-context_size:]
+    prompt = []
+    for idx, context in enumerate(reversed(contexts)):
+        if idx % 2 == 0:
+            prompt = [f"ASSISTANT: {context}"] + prompt
+        else:
+            prompt = [f"USER: {context}"] + prompt
+    prompt = "\n".join(prompt)
+    return prompt
+
+
+def interact_func(message, chat_history, max_context_size):
+    prompt = make_prompt(message, chat_history, max_context_size)
+    print(f"prompt: {prompt}")
+    generated = llm(prompt)
+    generated = generated.replace("\nUSER", "")
+    print(f"generated: {generated}")
+    chat_history.append((message, generated))
+    yield "", chat_history
+
+
+with gr.Blocks(theme="monochrome") as demo:
+    gr.Markdown(TITLE)
+    with gr.Accordion("Configs", open=False):
+        # max_context_size = the number of turns * 2
+        max_context_size = gr.Number(value=20, label="記憶する会話ターン数", precision=0)
+        max_length = gr.Number(value=128, label="最大文字数", precision=0)
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox()
+    clear = gr.Button("消す")
+    msg.submit(
+        make_prompt,
+        [msg, chatbot, max_context_size],
+        [msg, chatbot],
+        queue=False
+    ).then(
+        interact_func, chatbot, chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    demo.queue()
+    demo.launch(debug=True, share=True)