Update app.py

app.py

@@ -1,39 +1,58 @@
-import gradio as gr
 from huggingface_hub import InferenceClient
-import re
+import gradio as gr
+
+client = InferenceClient(
+    "mistralai/Mixtral-8x7B-Instruct-v0.1"
+)
 
-# Initialize the InferenceClient with the Mixtral model
-client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")
+def format_prompt(message, history):
+    prompt = "<s>"
+    for user_prompt, bot_response in history:
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
+    prompt += f"[INST] {message} [/INST]"
+    return prompt
 
-def translate_srt(file_info, target_language):
-    # file_info is a dictionary containing information about the file
-    # including its content
-    srt_content = file_info["content"].decode("utf-8")
-    lines = srt_content.split('\n')
+def generate_from_file(file_path, history, system_prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+    # Read the file content
+    with open(file_path, 'r', encoding='utf-8') as file:
+        file_content = file.read()
     
-    translated_lines = []
-    for line in lines:
-        if re.match(r"^\d+$", line) or re.match(r"^\d{2}:\d{2}:\d{2},\d{3} --> \d{2}:\d{2}:\d{2},\d{3}$", line):
-            translated_lines.append(line)  # Copy timestamps and numbers directly
-        elif line.strip() == "":
-            translated_lines.append(line)  # Preserve empty lines for formatting
-        else:
-            # Translate the text line
-            response = client(inputs={"inputs": line, "parameters": {"target_language": target_language}})
-            translated_lines.append(response[0]["generated_text"])
+    # You might need to modify this part to fit how you want to use the file content in your prompt
+    prompt = file_content[:1000]  # Example: using first 1000 characters of the file content
     
-    # Join the translated lines back into a single string
-    translated_srt_content = "\n".join(translated_lines)
-    return translated_srt_content
+    return generate(prompt, history, system_prompt, temperature, max_new_tokens, top_p, repetition_penalty)
+
+def generate(prompt, history, system_prompt, temperature=0.9, 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,
+    )
+
+    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    output = ""
+
+    for response in stream:
+        output += response.token.text
+        yield output
+    return output
 
-# Gradio interface
 iface = gr.Interface(
-    fn=translate_srt,
-    inputs=[gr.File(label="Upload SRT File"), gr.Dropdown(["fr", "en", "es", "de", "it", "pt"], label="Target Language")],
+    fn=generate_from_file,
+    inputs=[gr.File(label="Upload File"), gr.State(), gr.Textbox(label="System Prompt")],
     outputs="text",
-    title="SRT File Translator",
-    description="Translate SRT files to the selected language using Mixtral model."
+    title="SRT File Translation",
+    concurrency_limit=20,
 )
 
-# Launch the Gradio app
 iface.launch()