add inference logic and package requirements

app.py

@@ -1,7 +1,208 @@
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+import os
+import threading
+import time
+import argparse
+import logging
+from dataclasses import dataclass
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+import torch
+import sentencepiece as spm
+from transformers import GPTNeoXForCausalLM, GPTNeoXConfig
+from transformers.generation.streamers import BaseStreamer
+from huggingface_hub import hf_hub_download, login
+
+
+logger = logging.getLogger()
+logger.setLevel("INFO")
+
+gr_interface = None
+
+@dataclass
+class DefaultArgs:
+    hf_model_name_or_path: str = None
+    spm_model_path: str = None
+    env: str = "dev"
+    port: int = 7860
+    make_public: bool = False
+
+if os.getenv("RUNNING_ON_HF_SPACE"):
+    login(token=os.getenv("HF_TOKEN"))
+    hf_repo = "leemeng/stablelm-jp-alpha"
+    
+    args = DefaultArgs()
+    args.hf_model_name_or_path = hf_repo
+    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("--spm_model_path", type=str, required=True)   
+    parser.add_argument("--env", type=str, default="dev")      
+    parser.add_argument("--port", type=int, default=7860)      
+    parser.add_argument("--make_public", action='store_true') 
+    args = parser.parse_args()
+
+def load_model(
+    model_dir,
+):
+    config = GPTNeoXConfig.from_pretrained(model_dir)
+    config.is_decoder = True
+    model = GPTNeoXForCausalLM.from_pretrained(model_dir, config=config, torch_dtype=torch.bfloat16)
+    if torch.cuda.is_available():
+        model = model.to("cuda:0")
+    return model
+
+logging.info("Loading model")
+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
+        self.num_invoked = 0
+        self.prompt = ""
+        self.generated_text = ""
+        self.ended = False
+
+    
+    def put(self, t: torch.Tensor):
+        d = t.dim()
+        if d == 1:
+            pass
+        elif d == 2:
+            t = t[0]
+        else:
+            raise NotImplementedError
+        
+        t = [int(x) for x in t.numpy()]
+        
+        text = self.sp.decode_ids(t)
+        
+        if self.num_invoked == 0:
+            self.prompt = text
+            self.num_invoked += 1
+            return
+        
+        self.generated_text += text
+        # print(f"[streamer]: {self.generated_text}")
+        # yield text
+    
+    def end(self):
+        self.ended = True
+
+def user(user_message, history):
+    logging.debug(f"[user] user_message: {user_message}")
+    logging.debug(f"[user] history: {history}")
+    
+    res = ("", history + [[user_message, None]])
+    return res
+
+def bot(
+    history, 
+    temperature,
+    max_new_tokens,
+):
+    logging.debug(f"[bot] history: {history}")
+    logging.debug(f"temperature: {temperature}")
+    
+    # TODO: modify `<br>` back to `\n` based on the original user prinpt
+    prompt = history[-1][0]
+    
+    tokens = sp.encode(prompt)
+    input_ids = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(model.device)
+    
+    # TODO: parametrize setting on UI
+    do_sample = True
+        
+    streamer = SentencePieceStreamer(sp=sp)
+    
+    max_possilbe_new_tokens = model.config.max_position_embeddings - len(tokens)
+    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,
+        temperature=temperature,
+        max_new_tokens=max_possilbe_new_tokens,
+        do_sample=do_sample,
+        streamer=streamer,
+        # max_length=4096,
+        # top_k=100,
+        # top_p=0.9,
+        # repetition_penalty=1.0,
+        # num_return_sequences=2,
+        # num_beams=2,
+    ))
+    thr.start()
+    
+    history[-1][1] = ""
+    while not streamer.ended:
+        history[-1][1] = streamer.generated_text
+        time.sleep(0.05)
+        yield history
+    
+    # TODO: optimize for final few tokens
+    history[-1][1] = streamer.generated_text
+    yield history
+
+if gr_interface:
+    gr_interface.close(verbose=False)
+
+with gr.Blocks() as gr_interface:
+    chatbot = gr.Chatbot(label="StableLM JP Alpha").style(height=500)
+    
+    # generation params
+    with gr.Row():
+        temperature = gr.Slider(0, 1, value=0.7, step=0.05, label="Temperature")
+        max_new_tokens = gr.Slider(
+            128, 
+            model.config.max_position_embeddings, 
+            value=128, step=64, label="Max New Tokens")
+    
+    # prompt
+    # TODO: add more options
+#     prompt_options = gr.Dropdown(
+#         choices=[
+#             "運が良かったのか悪かったのか日本に帰ってきたタイミングでコロナが猛威を振るい始め、",
+#             """[問題]に対する[答え]を[選択肢]の中から選んでください。
+
+# [問題]: ある場所の周辺地域を指す言葉は？
+# [選択肢]: [空, オレゴン州, 街, 歩道橋, 近辺]
+# [答え]: 近辺
+
+# [問題]: 若くて世間に慣れていないことを何という？
+# [選択肢]: [青っぽい, 若い, ベテラン, 生々しい, 玄人]
+# [答え]: """
+#         ], 
+#         label="Prompt Options", 
+#         info="Select 1 option for quick start",
+#         allow_custom_value=False,
+#     )
+    prompt = gr.Textbox(label="Prompt", info="Pro tip: press Enter to submit directly")
+    
+    
+#     def on_prompt_options_change(pmt_opts, pmt):
+#         return pmt_opts
+    
+#     prompt_options.change(on_prompt_options_change, [prompt_options, prompt], prompt)
+
+    with gr.Row():
+        submit = gr.Button("Submit")
+        stop = gr.Button("Stop")
+    
+    clear = gr.Button("Clear History")
+    
+    # event handling
+    submit_event = prompt.submit(user, [prompt, chatbot], [prompt, chatbot], queue=False)\
+        .then(bot, [chatbot, temperature, max_new_tokens], chatbot, queue=True)
+    
+    submit_click_event = submit.click(user, [prompt, chatbot], [prompt, chatbot], queue=False)\
+        .then(bot, [chatbot, temperature, max_new_tokens], chatbot, queue=True)
+    
+    stop.click(None, None, None, cancels=[submit_event, submit_click_event], queue=False)
+    clear.click(lambda: None, None, chatbot, queue=False)
+
+gr_interface.queue()
+gr_interface.launch(server_port=args.port, share=args.make_public)   

requirements.txt

@@ -0,0 +1,6 @@
+gradio==3.28.1
+gradio_client==0.1.4
+torch==2.0.0
+sentencepiece==0.1.97
+transformers==4.28.1
+huggingface-hub==0.14.1