Update app.py

app.py

@@ -10,12 +10,40 @@ from transformers import (
 from datasets import load_dataset
 import logging
 import sys
+from urllib.parse import urlparse
 
 # Configure logging
 logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 
-def train(dataset_name: str, dataset_config: str = None):
+def parse_hf_dataset_url(url: str) -> tuple[str, str | None]:
+    """Parse Hugging Face dataset URL into (dataset_name, config)"""
+    parsed = urlparse(url)
+    path_parts = parsed.path.split('/')
+    
+    try:
+        # Find 'datasets' in path
+        datasets_idx = path_parts.index('datasets')
+    except ValueError:
+        raise ValueError("Invalid Hugging Face dataset URL")
+    
+    dataset_parts = path_parts[datasets_idx+1:]
+    dataset_name = "/".join(dataset_parts[0:2])
+    
+    # Try to find config (common pattern for datasets with viewer)
+    try:
+        viewer_idx = dataset_parts.index('viewer')
+        config = dataset_parts[viewer_idx+1] if viewer_idx+1 < len(dataset_parts) else None
+    except ValueError:
+        config = None
+    
+    return dataset_name, config
+
+def train(dataset_url: str):
     try:
+        # Parse dataset URL
+        dataset_name, dataset_config = parse_hf_dataset_url(dataset_url)
+        logging.info(f"Loading dataset: {dataset_name} (config: {dataset_config})")
+
         # Load model and tokenizer
         model_name = "microsoft/phi-2"
         tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
@@ -26,31 +54,45 @@ def train(dataset_name: str, dataset_config: str = None):
             tokenizer.pad_token = tokenizer.eos_token
 
         # Load dataset from Hugging Face Hub
-        logging.info(f"Loading dataset: {} (config: {dataset_config})")
         dataset = load_dataset(
             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
+            dataset_config,
+            trust_remote_code=True
         )
 
-        # Split into train/validation
-        dataset = dataset.train_test_split(test_size=0.1, seed=42)
+        # Handle dataset splits
+        if "train" not in dataset:
+            raise ValueError("Dataset must have a 'train' split")
+        
+        train_dataset = dataset["train"]
+        eval_dataset = dataset.get("validation", None)
+
+        # Split if no validation set
+        if eval_dataset is None:
+            split = train_dataset.train_test_split(test_size=0.1, seed=42)
+            train_dataset = split["train"]
+            eval_dataset = split["test"]
 
-        # Tokenization function (adjust based on dataset columns)
+        # Tokenization function
         def tokenize_function(examples):
             return tokenizer(
-                examples["text"],  # Replace "text" with your dataset's text column
+                examples["text"],  # Adjust column name as needed
                 padding="max_length",
                 truncation=True,
                 max_length=256,
                 return_tensors="pt",
             )
 
-        tokenized_dataset = dataset.map(
+        # Tokenize datasets
+        tokenized_train = train_dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=train_dataset.column_names
+        )
+        tokenized_eval = eval_dataset.map(
             tokenize_function,
             batched=True,
-            remove_columns=dataset["train"].column_names
+            remove_columns=eval_dataset.column_names
         )
 
         # Data collator
@@ -74,8 +116,8 @@ def train(dataset_name: str, dataset_config: str = None):
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=tokenized_dataset["train"],
-            eval_dataset=tokenized_dataset["test"],
+            train_dataset=tokenized_train,
+            eval_dataset=tokenized_eval,
             data_collator=data_collator,
         )
 
@@ -91,20 +133,22 @@ def train(dataset_name: str, dataset_config: str = None):
         logging.error(f"Training failed: {str(e)}")
         return f"❌ Training failed: {str(e)}"
 
-# Gradio UI with dataset input
+# Gradio UI with dataset URL input
 with gr.Blocks(title="Phi-2 Training") as demo:
     gr.Markdown("# 🚀 Train Phi-2 with HF Hub Data")
     
     with gr.Row():
-        dataset_name = gr.Textbox(label="Dataset Name", value="mozilla-foundation/common_voice_11_0")
-        dataset_config = gr.Textbox(label="Dataset Config (optional)", value="en")
+        dataset_url = gr.Textbox(
+            label="Dataset URL",
+            value="https://huggingface.co/datasets/mozilla-foundation/common_voice_11_0"
+        )
     
     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],
+        inputs=[dataset_url],
         outputs=status_output
     )