Update app.py

app.py

@@ -1,56 +1,83 @@
+import argparse
+import os
 import spaces
 
-import os
+import gradio as gr
+
 import json
-from vllm import LLM, SamplingParams
-from transformers import AutoTokenizer
+from threading import Thread
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+
+MAX_LENGTH = 4096
+DEFAULT_MAX_NEW_TOKENS = 1024
+
 
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base_model", type=str)  # model path
+    parser.add_argument("--n_gpus", type=int, default=1)  # n_gpu
+    return parser.parse_args()
 
 @spaces.GPU()
 def predict(message, history, system_prompt, temperature, max_tokens):
-    messages = [{"role": "system", "content": system_prompt}]
+    global model, tokenizer, device
+    messages = [{'role': 'system', 'content': system_prompt}]
     for human, assistant in history:
-        messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    stop_tokens = ["<|im_end|>", "<|endoftext|>", "<|im_start|>"]
-    sampling_params = SamplingParams(temperature=temperature, top_p=1, max_tokens=max_tokens, stop=stop_tokens)
-    completions = llm.generate(prompt, sampling_params)
-    for output in completions:
-        prompt = output.prompt
-        print('==========================question=============================')
-        print(prompt)
-        generated_text = output.outputs[0].text
-        print('===========================answer=============================')
-        print(generated_text)
-        for idx in range(len(generated_text)):
-                yield generated_text[:idx+1]
+        messages.append({'role': 'user', 'content': human})
+        messages.append({'role': 'assistant', 'content': assistant})
+    messages.append({'role': 'user', 'content': message})
+    problem = [tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)]
+    stop_tokens = ["<|endoftext|>", "<|im_end|>"]
+    streamer = TextIteratorStreamer(tokenizer, timeout=100.0, skip_prompt=True, skip_special_tokens=True)
+    enc = tokenizer(problem, return_tensors="pt", padding=True, truncation=True)
+    input_ids = enc.input_ids
+    attention_mask = enc.attention_mask
+
+    if input_ids.shape[1] > MAX_LENGTH:
+        input_ids = input_ids[:, -MAX_LENGTH:]
+
+    input_ids = input_ids.to(device)
+    attention_mask = attention_mask.to(device)
+    generate_kwargs = dict(
+        {"input_ids": input_ids, "attention_mask": attention_mask},
+        streamer=streamer,
+        do_sample=True,
+        top_p=0.95,
+        temperature=temperature,
+        max_new_tokens=DEFAULT_MAX_NEW_TOKENS,
+        use_cache=True,
+        eos_token_id=100278 # <|im_end|>
+    )
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        yield "".join(outputs)
+
 
 
 if __name__ == "__main__":
-    path = "stabilityai/stablelm-2-12b-chat"
-    tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
-    llm = LLM(model=path, tensor_parallel_size=1, trust_remote_code=True)
+    args = parse_args()
+    tokenizer = AutoTokenizer.from_pretrained("stabilityai/stablelm-2-chat", trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("stabilityai/stablelm-2-chat", trust_remote_code=True, torch_dtype=torch.bfloat16)
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
     gr.ChatInterface(
         predict,
-        title="LLM playground",
-        description="This is a LLM playground for StableLM",
+        title="StableLM 2 Chat - Demo",
+        description="StableLM 2 Chat - StabilityAI",
         theme="soft",
-        chatbot=gr.Chatbot(height=1400, label="Chat History",),
+        chatbot=gr.Chatbot(label="Chat History",),
         textbox=gr.Textbox(placeholder="input", container=False, scale=7),
         retry_btn=None,
         undo_btn="Delete Previous",
         clear_btn="Clear",
         additional_inputs=[
-            gr.Textbox("You are a hepful assistant.", label="System Prompt"),
-            gr.Slider(0, 1, 0.7, label="Temperature"),
+            gr.Textbox("You are a helpful assistant.", label="System Prompt"),
+            gr.Slider(0, 1, 0.5, label="Temperature"),
             gr.Slider(100, 2048, 1024, label="Max Tokens"),
         ],
         additional_inputs_accordion_name="Parameters",
-        examples=[
-            ["implement snake game using pygame"],
-            ["Can you explain briefly to me what is the Python programming language?"],
-            ["write a program to find the factorial of a number"],
-        ],
     ).queue().launch()