Update app.py

app.py

@@ -1,74 +1,60 @@
-import os
-import re
-import logging
 import gradio as gr
-import openai
+import re
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-print(os.environ)
-openai.api_base = os.environ.get("OPENAI_API_BASE")
-openai.api_key = os.environ.get("OPENAI_API_KEY")
+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)
 
-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."""
+BASE_SYSTEM_MESSAGE = "I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning."
 
 def make_prediction(prompt, max_tokens=None, temperature=None, top_p=None, top_k=None, repetition_penalty=None):
-    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"]
-
+    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
 
 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 []
-
-    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, ""
-
+    
+    # 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, ""
 
 start_message = ""
 
@@ -78,8 +64,6 @@ 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():