Create app.py

app.py

@@ -0,0 +1,99 @@
+import re
+
+import gradio as gr
+import torch
+from transformers import (AutoConfig, AutoModel, AutoModelForSeq2SeqLM,
+                          AutoTokenizer, LlamaForCausalLM, LlamaTokenizer)
+from vllm import LLM, SamplingParams
+
+model_id = "georgesung/llama2_7b_chat_uncensored"
+
+prompt_config = {
+    "system_header": None,
+    "system_footer": None,
+    "user_header": "### HUMAN:",
+    "user_footer": None,
+    "input_header": None,
+    "response_header": "### RESPONSE:",
+}
+
+def get_llm_response_chat(prompt):
+    outputs = llm.generate(prompt, sampling_params)
+    output = outputs[0].outputs[0].text
+    
+    # Remove trailing eos token
+    eos_token = llm.get_tokenizer().eos_token
+    if output.endswith(eos_token):
+        output = output[:-len(eos_token)]
+    return output
+
+def hist_to_prompt(history):
+    prompt = ""
+    if prompt_config["system_header"]:
+        system_footer = ""
+        if prompt_config["system_footer"]:
+            system_footer = prompt_config["system_footer"]
+        prompt += f"{prompt_config['system_header']}\n{SYSTEM_MESSAGE}{system_footer}\n\n"
+        
+    for i, (human_text, bot_text) in enumerate(history):
+        user_footer = ""
+        if prompt_config["user_footer"]:
+            user_footer = prompt_config["user_footer"]
+
+        prompt += f"{prompt_config['user_header']}\n{human_text}{user_footer}\n\n"
+
+        prompt += f"{prompt_config['response_header']}\n"
+        
+        if bot_text:
+            prompt += f"{bot_text}\n\n"
+    return prompt
+
+def get_bot_response(text):
+    bot_text_index = text.rfind(prompt_config['response_header'])
+    if bot_text_index != -1:
+        text = text[bot_text_index + len(prompt_config['response_header']):].strip()
+    return text
+
+def main():
+    # RE llama tokenizer:
+    # RuntimeError: Failed to load the tokenizer.
+    # If you are using a LLaMA-based model, use 'hf-internal-testing/llama-tokenizer' instead of the original tokenizer.
+    llm = LLM(model=model_id, tokenizer='hf-internal-testing/llama-tokenizer')
+
+    sampling_params = SamplingParams(temperature=0.01, top_p=0.1, top_k=40, max_tokens=2048)
+
+    tokenizer = llm.get_tokenizer()
+
+    with gr.Blocks() as demo:
+        gr.Markdown(
+        """
+        # Let's chat
+        """)
+        
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox()
+        clear = gr.Button("Clear")
+
+        def user(user_message, history):
+            return "", history + [[user_message, None]]
+
+        def bot(history):
+            hist_text = hist_to_prompt(history)
+
+            bot_message = get_llm_response_chat(hist_text) #+ tokenizer.eos_token
+            history[-1][1] = bot_message  # add bot message to overall history
+
+            return history
+
+        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+            bot, chatbot, chatbot
+        )
+        clear.click(lambda: None, None, chatbot, queue=False)
+
+    demo.queue()
+    demo.launch()
+
+
+
+if __name__ == "__main__":
+    main()