Delete train_model.py

train_model.py

@@ -1,172 +0,0 @@
-import os
-import glob
-import time
-import torch
-import logging
-from torch import nn, optim
-from accelerate import Accelerator
-from torch.utils.data import DataLoader
-from typing import Optional, Dict, List, Any
-from datasets import load_dataset, concatenate_datasets, Features, Value
-
-# Import your core model; choose one implementation for training.
-from model_Custm import Wildnerve_tlm01
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-# New helper function to flatten JSON with hierarchical markers.
-def flatten_json(data):
-    if isinstance(data, dict):
-        parts = []
-        for key, value in data.items():
-            parts.append(f"{key}:{{{flatten_json(value)}}}")
-        return " ".join(parts)
-    elif isinstance(data, list):
-        # Fixed the typo here: use "=" instead of "are"
-        parts = [flatten_json(item) for item in data]
-        return "[" + ", ".join(parts) + "]"
-    else:
-        return str(data)
-
-# New definition for convert_record, which uses flatten_json()
-def convert_record(record):
-    raw = record.get("text", "")
-    try:
-        import json
-        data = json.loads(raw)
-        combined = flatten_json(data)
-        return {"input": combined}
-    except Exception:
-        return {"input": raw}
-
-# Import tokenizer to convert text into tensor input
-from transformers import AutoTokenizer
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-# Updated get_dataset() function to load from Hugging Face repo
-def get_dataset(split="train", use_hf_data=True, dataset_repo="EvolphTech/data"):
-    if use_hf_data:
-        try:
-            logger.info(f"Loading dataset from Hugging Face: {dataset_repo}")
-            dataset = load_dataset(dataset_repo, split=split)
-            
-            # If the dataset has a 'text' column, use it directly
-            if 'text' in dataset.column_names:
-                dataset = dataset.map(lambda x: {"input": x["text"]})
-            else:
-                logger.warning(f"No 'text' column found in {dataset_repo}. Using first text column found.")
-                # Try to find a text column
-                text_columns = [col for col in dataset.column_names if dataset.features[col].dtype == 'string']
-                if text_columns:
-                    dataset = dataset.map(lambda x: {"input": x[text_columns[0]]})
-                else:
-                    raise ValueError(f"No text columns found in {dataset_repo}")
-                    
-            logger.info(f"Successfully loaded {len(dataset)} samples from Hugging Face")
-        except Exception as e:
-            logger.error(f"Failed to load dataset from Hugging Face: {e}")
-            logger.info("Falling back to local dataset")
-            return get_dataset(split=split, use_hf_data=False)
-    else:
-        # Fall back to the original local dataset loading logic
-        data_dir = r"c:\Users\User\OneDrive\Documents\tlm\Wildnerve-tlm_HF"
-        data_files = {
-            "train": os.path.join(data_dir, "train.json"),
-            "validation": os.path.join(data_dir, "validation.json")
-        }
-        features = Features({"text": Value("string")})
-        dataset = load_dataset("json", data_files=data_files, features=features, split=split, download_mode="force_redownload")
-        dataset = dataset.map(lambda x: {"input": x["text"]})
-        
-    class CustomDataset(torch.utils.data.Dataset):
-        def __init__(self, data):
-            self.data = data["input"]
-        def __len__(self):
-            return len(self.data)
-        def __getitem__(self, idx):
-            tokens = tokenizer(self.data[idx], truncation=True, padding="max_length", max_length=128, return_tensors="pt")
-            return tokens["input_ids"].squeeze(0)
-    
-    return CustomDataset(dataset)
-
-def train(use_hf_data=True, dataset_repo="EvolphTech/data"):
-    accelerator = Accelerator()
-    # Use the training split now
-    train_dataset = get_dataset("train", use_hf_data=use_hf_data, dataset_repo=dataset_repo)
-    train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
-
-    # Create your model (adjust constructor parameters as needed)
-    model = Wildnerve_tlm01(
-        vocab_size=30522,
-        specialization="general",
-        dataset_path="",
-        model_name="bert-base-uncased",
-        embedding_dim=256,
-        num_heads=4,
-        hidden_dim=256,
-        num_layers=2,
-        output_size=256,
-        dropout=0.1,
-        max_seq_length=128,
-        pooling_mode="mean",
-        use_pretrained_encoder=True
-    )
-    optimizer = optim.Adam(model.parameters(), lr=0.0001)
-    # Replace MSELoss with CrossEntropyLoss.
-    # Note: Assume model output logits are of shape [batch, seq_len, vocab_size]
-    criterion = nn.CrossEntropyLoss()
-    
-    model, optimizer, train_loader = accelerator.prepare(model, optimizer, train_loader)
-    
-    num_epochs = 50  # Change from 30 to 50
-    for epoch in range(num_epochs):
-        total_loss = 0.0
-        for batch in train_loader:
-            x = batch[..., :-1]  # omit last token for inputs
-            y = batch[..., 1:]   # omit first token for labels
-
-            optimizer.zero_grad()
-            output = model(x)  # shape is [batch_size, vocab_size]
-            
-            # Print shapes for debugging
-            logger.info(f"Epoch {epoch+1}, Output shape: {output.shape}, Target shape: {y.shape}")
-            
-            # Since the model returns logits for just one position, take the first token from y
-            # If your model really needs sequence data, you'd need a different handling strategy
-            target = y[:, 0].long()
-            
-            # Use target directly - no reshape needed since it's already 1D
-            loss = criterion(output, target)
-            
-            accelerator.backward(loss)
-            optimizer.step()
-            total_loss += loss.item()
-        avg_loss = total_loss / len(train_loader)
-        logger.info(f"Epoch {epoch+1}/{num_epochs}, Loss: {avg_loss:.4f}")
-        time.sleep(1)  # simulate longer training
-
-    # Save model weights as .pt then convert to .bin
-    results_dir = r"c:\Users\User\OneDrive\Documents\tlm\results"
-    os.makedirs(results_dir, exist_ok=True)
-    
-    pt_save_path = os.path.join(results_dir, "model_weights.pt")
-    torch.save(model.state_dict(), pt_save_path)
-    logger.info(f"Model weights saved to {pt_save_path}")
-    
-    # Convert .pt file to .bin (identical state_dict saved with .bin extension)
-    bin_save_path = os.path.join(results_dir, "model_weights.bin")
-    state_dict = torch.load(pt_save_path, weights_only=True)
-    torch.save(state_dict, bin_save_path)
-    logger.info(f"Model weights also saved as binary to {bin_save_path}")
-    
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser(description="Train the model")
-    parser.add_argument("--use_hf_data", action="store_true", help="Use data from Hugging Face repo")
-    parser.add_argument("--dataset_repo", type=str, default="EvolphTech/data", help="Hugging Face dataset repository")
-    parser.add_argument("--epochs", type=int, default=50, help="Number of training epochs")
-    
-    args = parser.parse_args()
-    
-    train(use_hf_data=args.use_hf_data, dataset_repo=args.dataset_repo)