update

app.py

@@ -1,35 +1,40 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, pipeline
+
 import gradio as gr
-import openai
 import os
 
 
-# Load openai key
-openai.api_key = os.getenv('OPENAI_KEY')
 
 # 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",
+    torch_dtype="auto",
+    device_map="auto"
+)
+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)
 
-# Create function to process prompt and append previous prompts as "context"
-def predict_prompt(input):
 
-    global message_history
-    message_history.append({"role": "user", "content": input})
-    create_prompt = openai.ChatCompletion.create(
-        model = "gpt-3.5-turbo",
-        messages = message_history
-    )
+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. """
+        return [
+            {"role": "system", "content": system_prompt},
+            {"role": "question", "content": prompt}
+        ]
 
-    reply_prompt = create_prompt.choices[0].message.content
-    # print(reply_prompt)
+    prompt = pipe.tokenizer.apply_chat_template(chat_template(prompt), 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()
 
-    # Append answer as assistant reply to keep history of prompts
-    message_history.append({"role": "assistant", "content": reply_prompt}) 
-    response = [(message_history[i]["content"], message_history[i+1]["content"]) for i in range(0, len(message_history) -1, 2)]
 
-    return response
 
-# Create UI using gradio
 with gr.Blocks(theme='abidlabs/dracula_test') as chatblock:
 
     gr.Markdown("<h1><center>Welcome to my personal AI assistant (powered by OpenAI)</center></h1>")
@@ -37,10 +42,10 @@ with gr.Blocks(theme='abidlabs/dracula_test') as chatblock:
     Chatbot = gr.Chatbot()
     with gr.Row():
         txt = gr.Textbox(
-            show_label=False, 
+            show_label=False,
             placeholder = initial_message).style(container=False)
         state = gr.State()
-        txt.submit(predict_prompt, txt, Chatbot)
+        txt.submit(chat, txt, Chatbot)
         txt.submit(None, None, txt, _js="() => {''}")
 
 chatblock.launch()

requirements.txt

@@ -1,2 +1,4 @@
 gradio==3.22.1
-openai==0.27.2
+accelerate
+bitsandbytes
+xformers