first commit

README.md

@@ -1,5 +1,5 @@
 ---
-title: Rinna Japanese Gpt Neox 3.6b Instruction Sft
+title: Chat with `rinna/japanese-gpt-neox-3.6b-instruction-sft`
 emoji: 👀
 colorFrom: gray
 colorTo: yellow

app.py

@@ -0,0 +1,71 @@
+import itertools
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import gradio as gr
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"device: {device}")
+
+tokenizer = AutoTokenizer.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", use_fast=False)
+model = AutoModelForCausalLM.from_pretrained("rinna/japanese-gpt-neox-3.6b-instruction-sft", device_map="auto", torch_dtype=torch.float16)
+model = model.to(device)
+
+
+@torch.no_grad()
+def inference_func(prompt, max_new_tokens=128, temperature=0.7):
+  token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+  output_ids = model.generate(
+      token_ids.to(model.device),
+      do_sample=True,
+      max_new_tokens=max_new_tokens,
+      temperature=temperature,
+      pad_token_id=tokenizer.pad_token_id,
+      bos_token_id=tokenizer.bos_token_id,
+      eos_token_id=tokenizer.eos_token_id
+  )
+  output = tokenizer.decode(output_ids.tolist()[0][token_ids.size(1):], skip_special_tokens=True)
+  output = output.replace("<NL>", "\n")
+  return output
+
+
+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"システム: {context}"] + prompt
+    else:
+      prompt = [f"ユーザー: {context}"] + prompt    
+  prompt = "<NL>".join(prompt)
+  return prompt
+
+def interact_func(message, chat_history, max_context_size, max_new_tokens, temperature):
+  prompt = make_prompt(message, chat_history, max_context_size)
+  print(f"prompt: {prompt}")
+  generated = inference_func(prompt, max_new_tokens, temperature)
+  print(f"generated: {generated}")
+  chat_history.append((message, generated))
+  return "", chat_history
+
+
+with gr.Blocks() as demo:
+  with gr.Accordion("Configs", open=False):
+      # max_context_size = the number of turns * 2
+      max_context_size = gr.Number(value=10, label="max_context_size", precision=0)
+      max_new_tokens = gr.Number(value=128, label="max_new_tokens", precision=0)
+      temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="temperature")
+  chatbot = gr.Chatbot()
+  msg = gr.Textbox()
+  clear = gr.Button("Clear")
+  msg.submit(interact_func, [msg, chatbot, max_context_size, max_new_tokens, temperature], [msg, chatbot])
+  clear.click(lambda: None, None, chatbot, queue=False)
+
+if __name__ == "__main__":
+    demo.launch(debug=True)

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+sentencepiece
+gradio
+accelerate>=0.12.0
+bitsandbytes>=0.31.5