import os
import modal

# Offload all the heavy dependency installations to the Modal cloud container
image = (
    modal.Image.debian_slim()
    .pip_install(
        "transformers",
        "datasets",
        "torch",
        "tokenizers",
        "huggingface_hub",
        "accelerate"
    )
)

app = modal.App("tralalabs-16m-qwen-master-pretrain")

@app.function(
    image=image,
    gpu="L40S",
    timeout=86400, # 24 hours max runtime allowed
    secrets=[modal.Secret.from_name("huggingface-secret")]
)
def train():
    import torch
    import torch.nn as nn
    from torch.utils.data import IterableDataset, DataLoader
    from datasets import load_dataset
    from tokenizers import Tokenizer, models, trainers, pre_tokenizers
    from transformers import PreTrainedTokenizerFast, Qwen2Config, Qwen2ForCausalLM
    from huggingface_hub import HfApi
    from torch.optim import AdamW
    
    print("Initialization started! Fetching data for Tokenizer and Training...")
    
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        print("Error: HF_TOKEN environment variable missing in your Modal secret.")
        return

    # 1. Stream the 85% / 10% / 5% mix
    try:
        ds_fw_2024 = load_dataset("HuggingFaceFW/fineweb-edu", "CC-MAIN-2024-18", split="train", streaming=True)
        ds_wiki = load_dataset("wikipedia", "20231101.en", split="train", streaming=True)
        ds_fw_2023 = load_dataset("HuggingFaceFW/fineweb-edu", "CC-MAIN-2023-50", split="train", streaming=True)
        
        def batch_iterator(batch_size=1000):
            fw_2024_iter = iter(ds_fw_2024)
            wiki_iter = iter(ds_wiki)
            fw_2023_iter = iter(ds_fw_2023)
            
            # Infinite loop generator for the massive 81k step pre-training run
            while True:
                batch = []
                for _ in range(int(batch_size * 0.85)): batch.append(next(fw_2024_iter)["text"])
                for _ in range(int(batch_size * 0.10)): batch.append(next(wiki_iter)["text"])
                for _ in range(int(batch_size * 0.05)): batch.append(next(fw_2023_iter)["text"])
                yield batch
    except Exception as e:
        print(f"Error setting up datasets: {e}")
        return

    # 2. Train Tokenizer (16k Vocab) using the first few batches
    print("Training 16k Byte-Level BPE Tokenizer...")
    raw_tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    raw_tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False)
    trainer = trainers.BpeTrainer(vocab_size=16000, special_tokens=["<unk>", "<s>", "</s>", "<pad>", "<mask>"])
    
    # Grab a finite chunk of data to train the vocabulary, then stop
    def tokenizer_iterator():
        iterator = batch_iterator(1000)
        for _ in range(20): 
            yield next(iterator)
            
    raw_tokenizer.train_from_iterator(tokenizer_iterator(), trainer=trainer)
    
    tokenizer = PreTrainedTokenizerFast(
        tokenizer_object=raw_tokenizer, 
        bos_token="<s>", 
        eos_token="</s>", 
        unk_token="<unk>", 
        pad_token="<pad>", 
        mask_token="<mask>"
    )
    tokenizer.pad_token = "<pad>"
    
    os.makedirs("./outputs", exist_ok=True)
    tokenizer.save_pretrained("./outputs")

    # 3. Model Hyperparameters: 16.7M params Qwen2 Architecture
    print("Configuring 16.7M Parameter Qwen2 Architecture...")
    config = Qwen2Config(
        vocab_size=16000,
        hidden_size=384,
        intermediate_size=1536,
        num_hidden_layers=6,
        num_attention_heads=6,
        num_key_value_heads=2, # GQA activated for maximum efficiency
        max_position_embeddings=1024,
        pad_token_id=3,
        bos_token_id=1,
        eos_token_id=2,
        tie_word_embeddings=True,
        rope_theta=10000.0
    )
    
    model = Qwen2ForCausalLM(config).to(device="cuda", dtype=torch.bfloat16)
    
    # 4. The 334M Token Training Loop
    print("Tokenizer baked! Starting massive gradient descent pre-training run...")
    optimizer = AdamW(model.parameters(), lr=6e-4, weight_decay=0.1)
    model.train()
    
    class ProportionalDataset(IterableDataset):
        def __init__(self, it): self.it = it
        def __iter__(self):
            for batch in self.it:
                for text in batch: yield text

    train_loader = DataLoader(ProportionalDataset(batch_iterator(batch_size=200)), batch_size=4)
    
    step = 0
    # 334,000,000 total tokens / (4 batch size * 1024 sequence length) = 81,543 steps
    TARGET_STEPS = 81543 
    
    for batch_text in train_loader:
        if step >= TARGET_STEPS: 
            break 
        
        optimizer.zero_grad()
        
        encodings = tokenizer(
            batch_text, 
            truncation=True, 
            max_length=1024, 
            padding="max_length", 
            return_tensors="pt"
        )
        
        input_ids = encodings["input_ids"].to("cuda")
        attention_mask = encodings["attention_mask"].to("cuda")
        
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=input_ids)
        loss = outputs.loss
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        
        # Log every 50 steps so the terminal doesn't get flooded
        if step % 50 == 0:
            print(f"Step {step}/{TARGET_STEPS} | Loss: {loss.item():.4f}")
            
        step += 1

    # 5. Save and Push the Final Master Weights
    print(f"Saving Final Learned Weights after {TARGET_STEPS} steps...")
    model.save_pretrained("./outputs")
    
    repo_id = "Tralalabs/TralaLabs-16M-Base"
    try:
        api = HfApi()
        api.create_repo(repo_id=repo_id, token=hf_token, exist_ok=True)
        api.upload_folder(folder_path="./outputs", repo_id=repo_id, repo_type="model", token=hf_token)
        print("Complete master success! Full 334M token Qwen model uploaded.")
    except Exception as e:
        print(f"Error uploading to HF: {e}")

@app.local_entrypoint()
def main():
    train.remote()