Update app.py

app.py

@@ -1,60 +1,74 @@
-import gradio as gr
+import os
 import re
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import logging
+import gradio as gr
+import openai
 
-model_name_or_path = "teknium/OpenHermes-2-Mistral-7B"
-model = AutoModelForCausalLM.from_pretrained(model_name_or_path,
-                                             device_map="auto",
-                                             trust_remote_code=False,
-                                             load_in_8bit=True,
-                                             revision="main")
-tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True)
+print(os.environ)
+openai.api_base = os.environ.get("OPENAI_API_BASE")
+openai.api_key = os.environ.get("OPENAI_API_KEY")
 
-BASE_SYSTEM_MESSAGE = "I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning."
+BASE_SYSTEM_MESSAGE = """I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning. 
+I am an assistant who thinks through their answers step-by-step to be sure I always get the right answer. 
+I think more clearly if I write out my thought process in a scratchpad manner first; therefore, I always explain background context, assumptions, and step-by-step thinking BEFORE trying to answer or solve anything."""
 
 def make_prediction(prompt, max_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    out = model.generate(input_ids, max_length=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty)
-    text = tokenizer.decode(out[0], skip_special_tokens=True)
-    yield text
+    completion = openai.Completion.create(model="teknium/OpenHermes-2-Mistral-7B", prompt=prompt, max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k, repetition_penalty=repetition_penalty, stream=True, stop=["</s>", "<|im_end|>"])
+    for chunk in completion:
+        yield chunk["choices"][0]["text"]
+
 
 def clear_chat(chat_history_state, chat_message):
     chat_history_state = []
     chat_message = ''
     return chat_history_state, chat_message
 
+
 def user(message, history):
     history = history or []
+    # Append the user's message to the conversation history
     history.append([message, ""])
     return "", history
 
+
 def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetition_penalty):
     history = history or []
-    
-    # A última mensagem do usuário
-    user_prompt = history[-1][0]
-    
-    # Definindo o template e o prompt
-    prompt_template = f'''system
-    {system_message.strip()}
-    user
-    {user_prompt}
-    assistant
-    '''
-    
-    # Preparando o input
-    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids  # .cuda() se você estiver usando GPU
-    
-    # Gerar a saída
-    output = model.generate(input_ids=input_ids, temperature=temperature, do_sample=True, top_p=top_p, top_k=top_k, max_length=max_tokens)
-    
-    # Decodificar a saída
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    # Atualizar o histórico
-    history[-1][1] += decoded_output
-    
-    yield history, history, ""
+
+    if system_message.strip():   
+        messages = "<|im_start|> "+"system\n" + system_message.strip() + "<|im_end|>\n" + \
+                   "\n".join(["\n".join(["<|im_start|> "+"user\n"+item[0]+"<|im_end|>", "<|im_start|> assistant\n"+item[1]+"<|im_end|>"])
+                        for item in history])
+    else:
+        messages = "<|im_start|> "+"system\n" + BASE_SYSTEM_MESSAGE + "<|im_end|>\n" + \
+                   "\n".join(["\n".join(["<|im_start|> "+"user\n"+item[0]+"<|im_end|>", "<|im_start|> assistant\n"+item[1]+"<|im_end|>"])
+                        for item in history])
+    # strip the last `<|end_of_turn|>` from the messages
+    messages = messages.rstrip("<|im_end|>")
+    # remove last space from assistant, some models output a ZWSP if you leave a space
+    messages = messages.rstrip()
+
+    # If temperature is set to 0, force Top P to 1 and Top K to -1
+    if temperature == 0:
+        top_p = 1
+        top_k = -1
+
+    prediction = make_prediction(
+        messages,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty,
+    )
+    for tokens in prediction:
+        tokens = re.findall(r'(.*?)(\s|$)', tokens)
+        for subtoken in tokens:
+            subtoken = "".join(subtoken)
+            answer = subtoken
+            history[-1][1] += answer
+            # stream the response
+            yield history, history, ""
+
 
 start_message = ""
 
@@ -64,6 +78,8 @@ CSS ="""
 #component-0 { height: 100%; }
 #chatbot { flex-grow: 1; overflow: auto; resize: vertical; }
 """
+
+#with gr.Blocks() as demo:
 with gr.Blocks(css=CSS) as demo:
     with gr.Row():
         with gr.Column():
@@ -109,4 +125,4 @@ with gr.Blocks(css=CSS) as demo:
     )
     stop.click(fn=None, inputs=None, outputs=None, cancels=[submit_click_event], queue=False)
 
-demo.queue(max_size=128, concurrency_count=48).launch(debug=True, server_name="0.0.0.0", server_port=7860)
+demo.queue(max_size=128, concurrency_count=48).launch(debug=True, server_name="0.0.0.0", server_port=7860)