Update app.py

app.py

@@ -34,13 +34,14 @@ The current release version of Breeze-7B is v1.0.
 """
 
 LICENSE = """
-
 """
 
 DEFAULT_SYSTEM_PROMPT = "You are a helpful AI assistant built by MediaTek Research. The user you are helping speaks Traditional Chinese and comes from Taiwan."
 
 API_URL = os.environ.get("API_URL")
 TOKEN = os.environ.get("TOKEN")
+TOKENIZER_REPO = "MediaTek-Research/Breeze-7B-Instruct-v1_0"
+API_MODEL_TYPE = "breeze-7b-instruct-v10"
 
 HEADERS = {
     "Authorization": f"Bearer {TOKEN}", 
@@ -48,34 +49,10 @@ HEADERS = {
     "accept": "application/json"
 }
 
-
 MAX_SEC = 30
 MAX_INPUT_LENGTH = 5000
 
-tokenizer = AutoTokenizer.from_pretrained("MediaTek-Research/Breeze-7B-Instruct-v0_1")
-
-def insert_to_db(prompt, response, temperature, top_p):
-    try:
-        #Establishing the connection
-        conn = psycopg2.connect(
-           database=os.environ.get("DB"), user=os.environ.get("USER"), password=os.environ.get("DB_PASS"), host=os.environ.get("DB_HOST"), port= '5432'
-        )
-        #Setting auto commit false
-        conn.autocommit = True
-        
-        #Creating a cursor object using the cursor() method
-        cursor = conn.cursor()
-    
-        # Preparing SQL queries to INSERT a record into the database.
-        cursor.execute(f"INSERT INTO breezedata(prompt, response, temperature, top_p) VALUES ('{prompt}', '{response}', {temperature}, {top_p})")
-    
-        # Commit your changes in the database
-        conn.commit()
-    
-        # Closing the connection
-        conn.close()
-    except:
-        pass
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_REPO, use_auth_token=os.environ.get("HF_TOKEN"))
 
 
 def refusal_condition(query):
@@ -105,19 +82,20 @@ with gr.Blocks() as demo:
     system_prompt = gr.Textbox(label='System prompt',
                                value=DEFAULT_SYSTEM_PROMPT,
                                lines=1)
-
+    
     with gr.Accordion(label='Advanced options', open=False):
+
         max_new_tokens = gr.Slider(
             label='Max new tokens',
             minimum=32,
-            maximum=1024,
+            maximum=2048,
             step=1,
-            value=512,
+            value=1024,
         )
         temperature = gr.Slider(
             label='Temperature',
             minimum=0.01,
-            maximum=1.0,
+            maximum=0.5,
             step=0.01,
             value=0.01,
         )
@@ -128,15 +106,8 @@ with gr.Blocks() as demo:
             step=0.01,
             value=0.01,
         )
-        repetition_penalty = gr.Slider(
-            label='Repetition Penalty',
-            minimum=0.1,
-            maximum=2,
-            step=0.01,
-            value=1.1,
-        )
 
-    chatbot = gr.Chatbot()
+    chatbot = gr.Chatbot(show_copy_button=True, show_share_button=True, )
     with gr.Row():
         msg = gr.Textbox(
             container=False,
@@ -157,7 +128,6 @@ with gr.Blocks() as demo:
 
     saved_input = gr.State()
 
-
     def user(user_message, history):
         return "", history + [[user_message, None]]
 
@@ -195,7 +165,7 @@ with gr.Blocks() as demo:
             # start_time = time.time()
 
     
-    def bot(history, max_new_tokens, temperature, top_p, system_prompt, repetition_penalty):
+    def bot(history, max_new_tokens, temperature, top_p, system_prompt):
         chat_data = []
         system_prompt = system_prompt.strip()
         if system_prompt:
@@ -217,19 +187,13 @@ with gr.Blocks() as demo:
             yield history
         else:
             data = {
-                "model_type": "breeze-7b-instruct-v10",
+                "model_type": API_MODEL_TYPE,
                 "prompt": str(message),
                 "parameters": {
                     "temperature": float(temperature),
                     "top_p": float(top_p),
                     "max_new_tokens": int(max_new_tokens),
-                    "repetition_penalty": float(repetition_penalty),
-
-                    "num_beams":1,    # w/o beam search
-                    "typical_p":0.99, 
-                    "top_k":61952,    # w/o top_k
-                    "do_sample": True,
-                    "min_length":1,
+                    "repetition_penalty": 1.1
                 }
             }
             
@@ -248,14 +212,13 @@ with gr.Blocks() as demo:
                 response = history[-1][1]
     
                 if refusal_condition(history[-1][1]):
-                    history[-1][1] = history[-1][1] + '\n\n**[免責聲明: Breeze-7B-Instruct 和 Breeze-7B-Instruct-64k 並未針對問答進行安全保護，因此語言模型的任何回應不代表 MediaTek Research 立場。]**'
+                    history[-1][1] = history[-1][1] + '\n\n**[免責聲明: 此模型並未針對問答進行安全保護，因此語言模型的任何回應不代表 MediaTek Research 立場。]**'
                     yield history
             else:
                 del history[-1]
                 yield history
 
         print('== Record ==\nQuery: {query}\nResponse: {response}'.format(query=repr(message), response=repr(history[-1][1])))
-        insert_to_db(message, response, float(temperature), float(top_p))
         
     msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
         fn=bot,
@@ -265,7 +228,6 @@ with gr.Blocks() as demo:
             temperature,
             top_p,
             system_prompt,
-            repetition_penalty,
         ],
         outputs=chatbot
     )
@@ -279,7 +241,6 @@ with gr.Blocks() as demo:
             temperature,
             top_p,
             system_prompt,
-            repetition_penalty,
         ],
         outputs=chatbot
     )
@@ -319,7 +280,6 @@ with gr.Blocks() as demo:
             temperature,
             top_p,
             system_prompt,
-            repetition_penalty,
         ],
         outputs=chatbot,
     )
@@ -342,5 +302,5 @@ with gr.Blocks() as demo:
 
     gr.Markdown(LICENSE)
 
-demo.queue(concurrency_count=2, max_size=128)
-demo.launch()
+demo.queue(concurrency_count=4, max_size=128)
+demo.launch()