Update app.py

app.py

@@ -630,90 +630,120 @@ def train_model(
         log.append(error_msg)
         return "\n".join(log)
     
-    # --- Download and Process Dataset ---
-    progress(0.4, desc="Downloading dataset...")
+    # --- Download and Load Dataset ---
+    progress(0.45, desc="Downloading dataset...")
+    log.append(f"Downloading dataset from {hf_dataset_repo_id}...")
     
     try:
-        dataset_path = "./downloaded_dataset_files"
+        # Download the dataset files
+        local_dataset_path = "./downloaded_dataset_files"
+        
+        # Correctly specify repo_type as "dataset"
         snapshot_download(
             repo_id=hf_dataset_repo_id,
-            local_dir=dataset_path,
-            use_auth_token=False,
+            local_dir=local_dataset_path,
+            repo_type="dataset",  # Important! Specifies this is a dataset repo
+            token=hf_token if hf_token and hf_token.strip() else None,  # Use token for auth
             resume_download=True
         )
-        log.append(f"Dataset repository content downloaded to: {dataset_path}")
-        
-        # Load dataset from PT files
-        progress(0.5, desc="Processing dataset...")
-        
-        # Load RVQ pairs
-        pair_files = glob.glob(f"{dataset_path}/*_rvq_pairs.pt")
-        log.append(f"Found {len(pair_files)} RVQ pair files.")
-        
-        all_pairs = []
-        for file in pair_files:
-            pairs = torch.load(file)
-            all_pairs.extend(pairs)
-        
-        log.append(f"Loaded a total of {len(all_pairs)} training pairs into memory.")
-        
-        # Process pairs into a format suitable for training
-        all_texts = []
-        for pair in all_pairs:
-            # Create instruction format
-            if isinstance(pair, dict):
-                instruction = pair.get("instruction", "")
-                input_text = pair.get("input", "")
-                output = pair.get("output", "")
+        log.append(f"Dataset files downloaded to {local_dataset_path}")
+        
+        # Check the structure of the downloaded files
+        log.append("Checking downloaded dataset structure...")
+        downloaded_files = glob.glob(f"{local_dataset_path}/**/*.pt", recursive=True)
+        log.append(f"Found {len(downloaded_files)} .pt files in the dataset directory")
+        
+        if len(downloaded_files) == 0:
+            log.append("No .pt files found. Checking for other file types...")
+            all_files = glob.glob(f"{local_dataset_path}/**/*.*", recursive=True)
+            log.append(f"All files found: {', '.join(all_files[:10])}")
+            if len(all_files) > 10:
+                log.append(f"...and {len(all_files) - 10} more files")
+        
+        # Look for the pairs directory
+        pairs_dir = os.path.join(local_dataset_path, "final_rvq_pairs")
+        if not os.path.exists(pairs_dir):
+            log.append(f"final_rvq_pairs directory not found. Looking for other possible directories...")
+            possible_dirs = [d for d in glob.glob(f"{local_dataset_path}/**/") if os.path.isdir(d)]
+            log.append(f"Available directories: {', '.join(possible_dirs)}")
+            
+            # Try to find any directory containing .pt files
+            for dir_path in possible_dirs:
+                if glob.glob(f"{dir_path}/*.pt"):
+                    pairs_dir = dir_path
+                    log.append(f"Using {pairs_dir} as the pairs directory.")
+                    break
+        
+        # If we found the pairs directory, we're good to go
+        if pairs_dir and os.path.exists(pairs_dir):
+            log.append(f"Using pairs directory: {pairs_dir}")
+            pt_files = glob.glob(f"{pairs_dir}/*.pt")
+            log.append(f"Found {len(pt_files)} .pt files in pairs directory")
+            
+            # Load the dataset from the files
+            progress(0.5, desc="Loading pairs from dataset files...")
+            log.append("Loading dataset pairs...")
+            
+            try:
+                # Load pairs from .pt files
+                pairs = []
+                for pt_file in tqdm(pt_files, desc="Loading .pt files"):
+                    pair_data = torch.load(pt_file)
+                    pairs.append(pair_data)
                 
-                # ALPACA format
-                if instruction and input_text:
-                    text = f"### Instruction: {instruction}\n### Input: {input_text}\n### Response: {output}"
-                elif instruction:
-                    text = f"### Instruction: {instruction}\n### Response: {output}"
-                else:
-                    text = output
-            else:
-                # Simple prompt-completion format
-                if isinstance(pair, tuple) and len(pair) == 2:
-                    prompt, completion = pair
-                    text = f"{prompt}{completion}"
+                log.append(f"Loaded {len(pairs)} conversation pairs")
+                
+                # Create a dataset from the pairs
+                dataset = Dataset.from_dict({
+                    "input_ids": [pair[0].tolist() for pair in pairs],
+                    "labels": [pair[1].tolist() for pair in pairs]
+                })
+                
+                # Split into training and validation sets
+                train_test_split = dataset.train_test_split(test_size=0.05)
+                train_dataset = train_test_split["train"]
+                
+                log.append(f"Created dataset with {len(train_dataset)} training examples")
+            
+            except Exception as e:
+                log.append(f"Error loading pair data: {e}")
+                
+                # Try an alternative approach - look for JSON or other formats
+                log.append("Attempting alternative dataset loading approaches...")
+                
+                # Search for JSON files
+                json_files = glob.glob(f"{local_dataset_path}/**/*.json", recursive=True)
+                if json_files:
+                    log.append(f"Found {len(json_files)} JSON files. Trying to load from these...")
+                    
+                    # Load from JSON
+                    combined_data = []
+                    for json_file in json_files[:5]:  # Start with a few files
+                        try:
+                            with open(json_file, 'r') as f:
+                                file_data = json.load(f)
+                                log.append(f"Successfully loaded {json_file}")
+                                # Print sample of the data structure
+                                log.append(f"Sample data structure: {str(file_data)[:500]}...")
+                                combined_data.append(file_data)
+                        except Exception as je:
+                            log.append(f"Error loading {json_file}: {je}")
+                    
+                    # If we loaded any data, try to create a dataset from it
+                    if combined_data:
+                        log.append("Attempting to create dataset from JSON data...")
+                        # This will need adapting based on the actual JSON structure
                 else:
-                    text = str(pair)
+                    log.append("No JSON files found. Looking for other formats...")
+                    # Add code for other formats if needed
+                
+                log.append("Failed to load dataset after multiple attempts.")
+                return "\n".join(log)
+                
+        else:
+            log.append("Could not locate pairs directory or any directory with .pt files.")
+            return "\n".join(log)
             
-            all_texts.append({"text": text})
-        
-        # Create HF dataset
-        train_dataset = Dataset.from_list(all_texts)
-        
-        # Function to tokenize the dataset
-        def tokenize_function(examples):
-            return tokenizer(
-                examples["text"],
-                padding=False,
-                truncation=True,
-                max_length=2048,
-                return_tensors=None,
-            )
-        
-        # Tokenize the dataset
-        tokenized_dataset = train_dataset.map(
-            tokenize_function,
-            batched=True,
-            remove_columns=["text"],
-            desc="Tokenizing dataset",
-        )
-        
-        train_dataset = tokenized_dataset
-        
-        # Data collator
-        from transformers import DataCollatorForLanguageModeling
-        
-        data_collator = DataCollatorForLanguageModeling(
-            tokenizer=tokenizer, 
-            mlm=False
-        )
-        
     except Exception as e:
         error_msg = f"Error loading dataset: {str(e)}"
         log.append(error_msg)