Update app.py

app.py

@@ -26,6 +26,10 @@ def translate_to_english(prompt):
     translated_prompt = GoogleTranslator(source='auto', target='en').translate(prompt)
     return translated_prompt
 
+def translate_to_persian(response):
+    translated_response = GoogleTranslator(source='auto', target='fa').translate(response)
+    return translated_response
+
 def load_models(inp):
     if VERBOSE == True:
         print(type(inp))
@@ -44,7 +48,7 @@ def format_prompt(message, history, cust_p):
     prompt += cust_p.replace("USER_INPUT", message)
     return prompt
 
-def chat_inf(system_prompt, prompt, history, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, cust_p):
+def chat_inf(system_prompt, prompt, history, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt, translate_to_persian):
     hist_len = 0
     client = clients[int(client_choice) - 1]
     if not history:
@@ -65,9 +69,9 @@ def chat_inf(system_prompt, prompt, history, memory, client_choice, seed, temp,
             max_new_tokens=tokens,
         )
         if system_prompt:
-            formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", memory[0 - chat_mem:], cust_p)
+            formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", memory[0 - chat_mem:], custom_prompt)
         else:
-            formatted_prompt = format_prompt(prompt, memory[0 - chat_mem:], cust_p)
+            formatted_prompt = format_prompt(prompt, memory[0 - chat_mem:], custom_prompt)
 
         translated_prompt = translate_to_english(formatted_prompt)
 
@@ -79,6 +83,8 @@ def chat_inf(system_prompt, prompt, history, memory, client_choice, seed, temp,
         output = ""
         for response in stream:
             output += response.token.text
+            if translate_to_persian:
+                output = translate_to_persian(output)
             yield [(prompt, output)], memory
         history.append((prompt, output))
         memory.append((prompt, output))
@@ -115,6 +121,7 @@ with gr.Blocks() as app:
                     custom_prompt = gr.Textbox(label="Modify Prompt Format", info="For testing purposes. 'USER_INPUT' is where 'SYSTEM_PROMPT, PROMPT' will be placed", lines=5, value="<start_of_turn>userUSER_INPUT<end_of_turn><start_of_turn>model")
             with gr.Column(scale=1):
                 with gr.Group():
+                    translate_to_persian_checkbox = gr.Checkbox(label="Translate to Persian", value=True)
                     rand = gr.Checkbox(label="Random Seed", value=True)
                     seed = gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, step=1, value=rand_val)
                     tokens = gr.Slider(label="Max new tokens", value=1600, minimum=0, maximum=8000, step=64, interactive=True, visible=True, info="The maximum number of tokens")
@@ -126,8 +133,8 @@ with gr.Blocks() as app:
     client_choice.change(load_models, client_choice, [chat_b])
     app.load(load_models, client_choice, [chat_b])
 
-    chat_sub = inp.submit(check_rand, [rand, seed], seed).then(chat_inf, [sys_inp, inp, chat_b, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt], [chat_b, memory])
-    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf, [sys_inp, inp, chat_b, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt], [chat_b, memory])
+    chat_sub = inp.submit(check_rand, [rand, seed], seed).then(chat_inf, [sys_inp, inp, chat_b, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt, translate_to_persian_checkbox], [chat_b, memory])
+    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf, [sys_inp, inp, chat_b, memory, client_choice, seed, temp, tokens, top_p, rep_p, chat_mem, custom_prompt, translate_to_persian_checkbox], [chat_b, memory])
 
     clear_btn.click(clear_fn, None, [inp, sys_inp, chat_b, memory])
 app.queue(default_concurrency_limit=10).launch()