Update app.py

app.py

@@ -2,15 +2,10 @@ from huggingface_hub import InferenceClient
 import gradio as gr
 from deep_translator import GoogleTranslator
 
+# Initialize the InferenceClient and the translators
 client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
-
-def translate_to_arabic(text):
-    translator = GoogleTranslator(source='auto', target='ar')
-    return translator.translate(text)
-
-def translate_to_english(text):
-    translator = GoogleTranslator(source='auto', target='en')
-    return translator.translate(text)
+translator_to_en = GoogleTranslator(source='arabic', target='english')
+translator_to_ar = GoogleTranslator(source='english', target='arabic')
 
 def format_prompt(message, history):
     prompt = "<s>"
@@ -20,39 +15,75 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
-def generate(prompt, history=[]):
+def generate(prompt, history, temperature=0.1, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+    temperature = float(temperature)
+    if temperature < 1e-2:
+        temperature = 1e-2
+    top_p = float(top_p)
+
+    generate_kwargs = dict(
+        temperature=temperature,
+        max_new_tokens=max_new_tokens,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
+        do_sample=True,
+        seed=42,
+    )
+
     # Translate the Arabic prompt to English
-    prompt_in_english = translate_to_english(prompt)
-    # Call the format_prompt function to format the input for the model
-    formatted_prompt = format_prompt(prompt_in_english, history)
-    
-    generate_kwargs = {
-        "temperature": 0.1,
-        "max_new_tokens": 256,
-        "top_p": 0.95,
-        "repetition_penalty": 1.0,
-        "do_sample": True,
-        "seed": 42,  # Seed for reproducibility, remove or change if randomness is preferred
-    }
-
-    # Generate the response from the model
+    translated_prompt = translator_to_en.translate(prompt)
+    formatted_prompt = format_prompt(translated_prompt, history)
+
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    
     output = ""
-    for response in stream:
-        if hasattr(response, 'text'):  # Checks if the 'text' attribute is present
-            output += response.text
 
-    # Translate the English response back to Arabic
-    response_in_arabic = translate_to_arabic(output)
-    return response_in_arabic
+    for response in stream:
+        output += response.token.text
+        yield translator_to_ar.translate(output)  # Translate the response back to Arabic
+    return output
 
+additional_inputs=[
+    gr.Slider(
+        label="Temperature",
+        value=0.9,
+        minimum=0.0,
+        maximum=1.0,
+        step=0.05,
+        interactive=True,
+        info="Higher values produce more diverse outputs",
+    ),
+    gr.Slider(
+        label="Max new tokens",
+        value=256,
+        minimum=0,
+        maximum=1048,
+        step=64,
+        interactive=True,
+        info="The maximum numbers of new tokens",
+    ),
+    gr.Slider(
+        label="Top-p (nucleus sampling)",
+        value=0.90,
+        minimum=0.0,
+        maximum=1,
+        step=0.05,
+        interactive=True,
+        info="Higher values sample more low-probability tokens",
+    ),
+    gr.Slider(
+        label="Repetition penalty",
+        value=1.2,
+        minimum=1.0,
+        maximum=2.0,
+        step=0.05,
+        interactive=True,
+        info="Penalize repeated tokens",
+    )
+]
 
-iface = gr.Interface(
+gr.ChatInterface(
     fn=generate,
-    inputs=[gr.Textbox(lines=5, placeholder='Type your Arabic query here...', label='Arabic Query')],
-    outputs='text',
-    title="DorjGPT Arabic-English Translation Chatbot"
-)
-
-iface.launch()
+    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
+    additional_inputs=additional_inputs,
+    title="DorjGPT interface with Arabic Translation"
+).launch(show_api=True)