update version with fujiki-tuned model

app.py

@@ -10,6 +10,7 @@ from dataclasses import dataclass
 
 import torch
 import sentencepiece as spm
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import GPTNeoXForCausalLM, GPTNeoXConfig
 from transformers.generation.streamers import BaseStreamer
 from huggingface_hub import hf_hub_download, login
@@ -20,11 +21,12 @@ logger.setLevel("INFO")
 
 gr_interface = None
 
-VERSION = "0.0.0-a.1"
+VERSION = "0.1.0-a.1"
 
 @dataclass
 class DefaultArgs:
     hf_model_name_or_path: str = None
+    hf_tokenizer_name_or_path: str = None
     spm_model_path: str = None
     env: str = "dev"
     port: int = 7860
@@ -35,10 +37,13 @@ if os.getenv("RUNNING_ON_HF_SPACE"):
     hf_repo = os.getenv("HF_MODEL_REPO")
     args = DefaultArgs()
     args.hf_model_name_or_path = hf_repo
+    args.hf_tokenizer_name_or_path = os.path.join(hf_repo, "tokenizer")
     args.spm_model_path = hf_hub_download(repo_id=hf_repo, filename="sentencepiece.model")
+    
 else:
     parser = argparse.ArgumentParser(description="")
     parser.add_argument("--hf_model_name_or_path", type=str, required=True)
+    parser.add_argument("--hf_tokenizer_name_or_path", type=str, required=False)
     parser.add_argument("--spm_model_path", type=str, required=True)   
     parser.add_argument("--env", type=str, default="dev")      
     parser.add_argument("--port", type=int, default=7860)      
@@ -60,9 +65,15 @@ model = load_model(args.hf_model_name_or_path)
 sp = spm.SentencePieceProcessor(model_file=args.spm_model_path)
 logging.info("Finished loading model")
 
-class SentencePieceStreamer(BaseStreamer):
-    def __init__(self, sp: spm.SentencePieceProcessor):
-        self.sp = sp
+tokenizer = AutoTokenizer.from_pretrained(
+    args.hf_model_name_or_path, 
+    subfolder="tokenizer",
+    use_fast=False
+)
+
+class TokenizerStreamer(BaseStreamer):
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
         self.num_invoked = 0
         self.prompt = ""
         self.generated_text = ""
@@ -80,7 +91,10 @@ class SentencePieceStreamer(BaseStreamer):
         
         t = [int(x) for x in t.numpy()]
         
-        text = self.sp.decode_ids(t)
+        text = tokenizer.decode(t)
+        if text in [tokenizer.bos_token, tokenizer.eos_token]:
+            text = ""
+        
         
         if self.num_invoked == 0:
             self.prompt = text
@@ -93,6 +107,35 @@ class SentencePieceStreamer(BaseStreamer):
     def end(self):
         self.ended = True
 
+INPUT_PROMPT = """以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。
+
+### 指示: 
+{instruction}
+
+### 入力: 
+{input}
+
+### 応答: """
+
+NO_INPUT_PROMPT = """以下は、タスクを説明する指示と、文脈のある入力の組み合わせです。要求を適切に満たす応答を書きなさい。
+
+### 指示: 
+{instruction}
+
+### 応答: """
+
+
+def postprocess_output(output):
+    output = output\
+        .split('### 応答:')[1]\
+        .split('###')[0]\
+        .split('##')[0]\
+        .lstrip(tokenizer.bos_token)\
+        .rstrip(tokenizer.eos_token)\
+        .replace("###", "")\
+        .strip()
+    return output
+
 def generate(
     prompt,
     max_new_tokens,
@@ -104,29 +147,29 @@ def generate(
 ):
     log = dict(locals())
     logging.debug(log)
-    
-    tokens = sp.encode(prompt)
-    input_ids = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(model.device)
-        
-    streamer = SentencePieceStreamer(sp=sp)
-    
-    max_possilbe_new_tokens = model.config.max_position_embeddings - len(tokens)
+
+    input_text = NO_INPUT_PROMPT.format(instruction=prompt)
+    input_ids = tokenizer.encode(input_text, add_special_tokens=False, return_tensors="pt")
+
+    streamer = TokenizerStreamer(tokenizer=tokenizer)
+
+    max_possilbe_new_tokens = model.config.max_position_embeddings - input_ids.shape[0]
     max_possilbe_new_tokens = min(max_possilbe_new_tokens, max_new_tokens)
     
     thr = threading.Thread(target=model.generate, args=(), kwargs=dict(
-        input_ids=input_ids,
+        input_ids=input_ids.to(model.device),
         do_sample=do_sample,
         temperature=temperature,
         repetition_penalty=repetition_penalty,
         no_repeat_ngram_size=no_repeat_ngram_size,
         max_new_tokens=max_possilbe_new_tokens,
+        pad_token_id=tokenizer.pad_token_id,
+        bos_token_id=tokenizer.bos_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        bad_words_ids=[[tokenizer.unk_token_id]],
         streamer=streamer,
-        # max_length=4096,
-        # top_k=100,
-        # top_p=0.9,
-        # num_return_sequences=2,
-        # num_beams=2,
     ))
+
     thr.start()
     
     while not streamer.ended:
@@ -165,10 +208,10 @@ if gr_interface:
 
 with gr.Blocks() as gr_interface:
     with gr.Row():
-        gr.Markdown(f"# 日本語 StableLM Pre-Alpha ({VERSION})")
+        gr.Markdown(f"# 日本語 StableLM Tuned Pre-Alpha ({VERSION})")
         # gr.Markdown(f"バージョン：{VERSION}")
     with gr.Row():
-        gr.Markdown("この言語モデルは Stability AI Japan が開発した初期バージョンの日本語モデルです。このモデルは質問応答に特化していません。そのため、期待する回答がプロンプトの自然な続きとなるようにプロンプトを設定する必要があります。例を挙げると「人生の意味は何ですか？」ではなく、「私が思った人生の意味は」のようにプロンプトを設定してください。")
+        gr.Markdown("この言語モデルは Stability AI Japan が開発した初期バージョンの日本語モデルです。モデルは「プロンプト」に入力した聞きたいことに対して、それらしい応答をすることができます。")
     with gr.Row():
         
         # left panel
@@ -180,7 +223,7 @@ with gr.Blocks() as gr_interface:
                 
                 # hidden default params
                 do_sample = gr.Checkbox(True, label="Do Sample", info="サンプリング生成", visible=True)
-                no_repeat_ngram_size = gr.Slider(0, 10, value=5, step=1, label="No Repeat Ngram Size", visible=False)
+                no_repeat_ngram_size = gr.Slider(0, 10, value=3, step=1, label="No Repeat Ngram Size", visible=False)
                 
                 # visible params
                 max_new_tokens = gr.Slider(
@@ -192,7 +235,7 @@ with gr.Blocks() as gr_interface:
                     info="生成するトークンの最大数を指定する",
                 )
                 temperature = gr.Slider(
-                    0, 1, value=0.7, step=0.05, label="temperature", 
+                    0, 1, value=0.1, step=0.05, label="temperature", 
                     info="低い値は出力をより集中させて決定論的にする")
                 repetition_penalty = gr.Slider(
                     1, 1.5, value=1.2, step=0.05, label="frequency penalty",
@@ -214,7 +257,7 @@ with gr.Blocks() as gr_interface:
             with gr.Box():
                 textbox_prompt = gr.Textbox(
                     label="プロンプト",
-                    placeholder="私が思った人生の意味は",
+                    placeholder="日本の首都は？",
                     interactive=True,
                     lines=5,
                     value=""