Update app.py

app.py

@@ -14,7 +14,7 @@ import sys
 # Configure logging
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
-def train():
+def train(dataset_name: str, dataset_config: str = None):
     try:
         # Load model and tokenizer
         model_name = "microsoft/phi-2"
@@ -25,19 +25,22 @@ def train():
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
 
-        # Load dataset
+        # Load dataset from Hugging Face Hub
+        logging.info(f"Loading dataset: {eswardivi/medical_qa} (config: {dataset_config})")
         dataset = load_dataset(
-            "csv",
-            data_files={
-                "train": "data/train/data.csv",
-                "validation": "data/validation/data.csv"
-            }
+            dataset_name,
+            dataset_config,  # Optional config (e.g., language for Common Voice)
+            split="train+validation",  # Combine splits
+            trust_remote_code=True  # Required for some datasets
         )
 
-        # Tokenization function
+        # Split into train/validation
+        dataset = dataset.train_test_split(test_size=0.1, seed=42)
+
+        # Tokenization function (adjust based on dataset columns)
         def tokenize_function(examples):
             return tokenizer(
-                examples["text"],
+                examples["text"],  # Replace "text" with your dataset's text column
                 padding="max_length",
                 truncation=True,
                 max_length=256,
@@ -47,7 +50,7 @@ def train():
         tokenized_dataset = dataset.map(
             tokenize_function,
             batched=True,
-            remove_columns=["text"]
+            remove_columns=dataset["train"].column_names
         )
 
         # Data collator
@@ -72,7 +75,7 @@ def train():
             model=model,
             args=training_args,
             train_dataset=tokenized_dataset["train"],
-            eval_dataset=tokenized_dataset["validation"],
+            eval_dataset=tokenized_dataset["test"],
             data_collator=data_collator,
         )
 
@@ -88,16 +91,20 @@ def train():
         logging.error(f"Training failed: {str(e)}")
         return f"❌ Training failed: {str(e)}"
 
-# Gradio UI
+# Gradio UI with dataset input
 with gr.Blocks(title="Phi-2 Training") as demo:
-    gr.Markdown("# 🚀 Train Phi-2 on CPU")
+    gr.Markdown("# 🚀 Train Phi-2 with HF Hub Data")
     
     with gr.Row():
-        start_btn = gr.Button("Start Training", variant="primary")
-        status_output = gr.Textbox(label="Status", interactive=False)
+        dataset_name = gr.Textbox(label="Dataset Name", value="mozilla-foundation/common_voice_11_0")
+        dataset_config = gr.Textbox(label="Dataset Config (optional)", value="en")
+    
+    start_btn = gr.Button("Start Training", variant="primary")
+    status_output = gr.Textbox(label="Status", interactive=False)
     
     start_btn.click(
         fn=train,
+        inputs=[dataset_name, dataset_config],
         outputs=status_output
     )