Update app.py

app.py

@@ -1,83 +1,56 @@
-import argparse
-import os
 import spaces
 
-import gradio as gr
-
+import os
 import json
-from threading import Thread
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-
-MAX_LENGTH = 4096
-DEFAULT_MAX_NEW_TOKENS = 1024
-
+from vllm import LLM, SamplingParams
+from transformers import AutoTokenizer
 
-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):
-    global model, tokenizer, device
-    messages = [{'role': 'system', 'content': system_prompt}]
+    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})
-    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=0.5,
-        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)
-
+        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]
 
 
 if __name__ == "__main__":
-    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)
+    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)
     gr.ChatInterface(
         predict,
-        title="StableLM 2 Chat - Demo",
-        description="StableLM 2 Chat - StabilityAI",
+        title="LLM playground",
+        description="This is a LLM playground for StableLM",
         theme="soft",
-        chatbot=gr.Chatbot(label="Chat History",),
+        chatbot=gr.Chatbot(height=1400, 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 helpful assistant.", label="System Prompt"),
-            gr.Slider(0, 1, 0.5, label="Temperature"),
+            gr.Textbox("You are a hepful assistant.", label="System Prompt"),
+            gr.Slider(0, 1, 0.7, 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()