Update app.py

app.py

@@ -512,7 +512,8 @@ def translate_dataset(train_url, local_parquet_path, input_file_path, output_fil
 logging.basicConfig(level=logging.INFO, filename='translation.log', filemode='a',
                             format='%(asctime)s - %(levelname)s - %(message)s')
 
-def main(model_id, dataset_url, model_type, output_dataset_name):
+# Main function to handle the translation workflow
+def main(model_id, dataset_url, model_type, output_dataset_name, range_specification):
     try:
         # Login to Hugging Face
         token = login()
@@ -543,7 +544,7 @@ def main(model_id, dataset_url, model_type, output_dataset_name):
                 "input_file": f"{model_type}_en.jsonl",
                 "output_file": f"{model_type}_de.jsonl",
                 "raw_file": f"{model_type}_de_raw.jsonl",
-                "range_spec": "1-",
+                "range_spec": range_specification,
                 "model_type": model_type
             }
 
@@ -589,6 +590,7 @@ with gr.Blocks() as demo:
             dataset_url = gr.Textbox(label="Dataset URL", lines=1)
             model_type = gr.Dropdown(choices=["mix", "orpo", "ufb"], label="Model Type")
             output_dataset_name = gr.Textbox(label="Output Dataset Name", lines=1)
+            range_specification = gr.Textbox(label="Range Specification", lines=1, placeholder="e.g., 1-100, 200-250")
             login_button = gr.Button("Login to Hugging Face")
 
         with gr.Column():
@@ -596,6 +598,6 @@ with gr.Blocks() as demo:
             logout_button = gr.Button("Logout")
 
     submit_btn = gr.Button("Translate Dataset", variant="primary")
-    submit_btn.click(main, inputs=[model_id, dataset_url, model_type, output_dataset_name], outputs=output)
+    submit_btn.click(main, inputs=[model_id, dataset_url, model_type, output_dataset_name, range_specification], outputs=output)
 
 demo.launch()