Update app.py

app.py

@@ -5,15 +5,15 @@ client = InferenceClient(
     "mistralai/Mixtral-8x7B-Instruct-v0.1"
 )
 
-def format_prompt(message, history):
+def format_prompt(message, history, target_language):
     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]"
+    prompt += f"[INST] {message} [/INST] [Translate to: {target_language}]"
     return prompt
 
-def generate_from_srt(file_content, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
+def generate_from_srt(file_content, target_language, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition_penalty=1.0):
     # Process the SRT file content as needed before using it as a prompt
     # For example, extracting text and removing timestamps if necessary
     # Directly using the file content for simplicity here
@@ -32,7 +32,7 @@ def generate_from_srt(file_content, temperature=0.9, max_new_tokens=256, top_p=0
         seed=42,
     )
 
-    formatted_prompt = format_prompt(file_content, [])
+    formatted_prompt = format_prompt(file_content, [], target_language)
     stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
     output = ""
 
@@ -40,7 +40,7 @@ def generate_from_srt(file_content, temperature=0.9, max_new_tokens=256, top_p=0
         output += response.token.text
     return output
 
-def handle_file(file_info):
+def handle_file(file_info, target_language):
     # Directly use the file content if it's a text file
     if isinstance(file_info, str):
         file_content = file_info
@@ -48,14 +48,15 @@ def handle_file(file_info):
         # If file_info is not a string, it might be a binary file
         file_content = file_info.decode('utf-8')
     
-    return generate_from_srt(file_content)
+    return generate_from_srt(file_content, target_language)
 
 iface = gr.Interface(
     fn=handle_file,
-    inputs=gr.File(label="Upload SRT File"),
+    inputs=[gr.File(label="Upload SRT File"), gr.Textbox(label="Target Language", placeholder="Enter target language (e.g., Spanish, French)")],
     outputs="text",
     title="SRT File Translation",
+    description="Translate the content of SRT files to the specified language.",
     concurrency_limit=20,
 )
 
-iface.launch()
+iface.launch()