Update app.py

app.py

@@ -1,51 +1,40 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import AutoTokenizer, pipeline
-
+import json
+import random
 import gradio as gr
-import os
 
+from tqdm import tqdm
 
+from datasets import load_dataset
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    TextStreamer
+)
 
-# Initialize message history array
-message_history = []
-initial_message = "Please write your prompt here and press 'enter'"
-device = "cuda" # the device to load the model onto
 
 model = AutoModelForCausalLM.from_pretrained(
-    "DuongTrongChi/Rikka-1.8B-v2.2",
+    'DuongTrongChi/Rikka-1.8B-v2.2',
     torch_dtype="auto",
     device_map="auto"
 )
-tokenizer = AutoTokenizer.from_pretrained("DuongTrongChi/Rikka-1.8B-v2.2")
+
+tokenizer = AutoTokenizer.from_pretrained('DuongTrongChi/Rikka-1.8B-v2.2')
 eos_token = tokenizer("<|im_end|>",add_special_tokens=False)["input_ids"][0]
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 
-def chat(prompt):
-    def chat_template(prompt):
-        system_prompt = """Bạn là một trợ lý hữu ích, tôn trọng và trung thực. Luôn trả lời một cách hữu ích nhất có thể trong khi vẫn an toàn. Câu trả lời của bạn không được bao gồm bất kỳ nội dung có hại, phi đạo đức, phân biệt chủng tộc, phân biệt giới tính, độc hại, nguy hiểm hoặc bất hợp pháp. """
+def chat(message, history=None):
+    def chat_template(message):
+        system_prompt = """Bạn là một trợ lý hữu ích, tôn trọng và trung thực. Luôn trả lời một cách hữu ích nhất có thể trong khi vẫn an toàn. """
         return [
             {"role": "system", "content": system_prompt},
-            {"role": "question", "content": prompt}
+            {"role": "question", "content": message}
         ]
-
-    prompt = pipe.tokenizer.apply_chat_template(chat_template(prompt), tokenize=False, add_generation_prompt=True)
+    prompt = pipe.tokenizer.apply_chat_template(chat_template(message), tokenize=False, add_generation_prompt=True)
     outputs = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.3, top_k=50, top_p=0.95, eos_token_id=eos_token, pad_token_id=eos_token)
-    return outputs[0]['generated_text'][len(prompt):].strip()
 
+    return outputs[0]['generated_text'][len(prompt):].strip()
 
 
-with gr.Blocks(theme='abidlabs/dracula_test') as chatblock:
-
-    gr.Markdown("<h1><center>Welcome to my personal AI assistant (powered by OpenAI)</center></h1>")
-
-    Chatbot = gr.Chatbot()
-    with gr.Row():
-        txt = gr.Textbox(
-            show_label=False,
-            placeholder = initial_message).style(container=False)
-        state = gr.State()
-        txt.submit(chat, txt, Chatbot)
-        txt.submit(None, None, txt, _js="() => {''}")
-
-chatblock.launch()
+gr.ChatInterface(chat).launch()