tadc.py

import json
import os
import re
import time
from datetime import datetime

import torch
from datasets import Dataset
from transformers import (
    Trainer,
    TrainerCallback,
    TrainingArguments,
    GPT2Config,
    GPT2LMHeadModel,
    GPT2TokenizerFast,
    LlamaConfig,
    LlamaForCausalLM,
    LlamaTokenizer,
)

device = "cuda" if torch.cuda.is_available() else "cpu"

print("=" * 60)
print("TADC TRAINER")
print("=" * 60)
print(f"Using device: {device}")

if torch.cuda.is_available():
    print(f"GPU: {torch.cuda.get_device_name(0)}")
    print(f"CUDA: {torch.version.cuda}")

OUTPUT_DIR = input("Enter the path where you want to save the model: ").strip()
epoch_num = int(input("Number of epochs: ").strip())

DATASET_PATH = "dataset-tadc.txt"
MAX_LENGTH = 4096
CHUNK_OVERLAP = 256

TRAINING_WINDOW_ENABLED = True
TRAINING_WINDOW_START = 5
TRAINING_WINDOW_END = 18

SIZE_MODES = {
    "banana": dict(n_layer=2, n_head=2, n_embd=32, hidden_size=32, intermediate_size=128),
    "nano": dict(n_layer=2, n_head=2, n_embd=64, hidden_size=64, intermediate_size=256),
    "small": dict(n_layer=4, n_head=4, n_embd=128, hidden_size=128, intermediate_size=512),
    "medium": dict(n_layer=6, n_head=6, n_embd=384, hidden_size=384, intermediate_size=1536),
    "large": dict(n_layer=12, n_head=12, n_embd=768, hidden_size=768, intermediate_size=3072),
    "larger": dict(n_layer=24, n_head=16, n_embd=1024, hidden_size=1024, intermediate_size=4096),
}

ROLE_PREFIX = "<|role|>"
END_TOKEN = "<|end|>"

def is_within_training_window():
    if not TRAINING_WINDOW_ENABLED:
        return True
    hour = datetime.now().hour
    return TRAINING_WINDOW_START <= hour < TRAINING_WINDOW_END

def wait_for_training_window():
    if is_within_training_window():
        return
    now = datetime.now()
    print(
        f"\nOutside training window ({now.strftime('%H:%M')}). "
        f"Pausing until {TRAINING_WINDOW_START:02d}:00..."
    )
    while not is_within_training_window():
        time.sleep(60)
    print(f"Training window open ({datetime.now().strftime('%H:%M')}). Resuming...")

class TimeWindowCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        if TRAINING_WINDOW_ENABLED and not is_within_training_window():
            wait_for_training_window()


def parse_tadc_turns(text):
    """
    Parse turns from text that looks like:
        <|role|>caine
        hello
        <|role|>pomni
        where am i?
    """
    lines = text.splitlines()
    turns = []
    current_role = None
    current_message = []

    def flush():
        nonlocal current_role, current_message
        if current_role is None:
            return
        message = "\n".join(current_message).strip()
        if message:
            turns.append((current_role, message))
        current_role = None
        current_message = []

    i = 0
    while i < len(lines):
        line = lines[i].rstrip("\n")
        stripped = line.strip()

        if not stripped:
            i += 1
            continue

        if stripped.startswith(ROLE_PREFIX):
            flush()
            remainder = stripped[len(ROLE_PREFIX):].strip()
            if remainder:
                current_role = remainder
                i += 1
                continue
            if i + 1 < len(lines):
                next_line = lines[i + 1].strip()
                if next_line:
                    current_role = next_line
                    i += 2
                    continue

        if current_role is not None:
            current_message.append(line)

        i += 1

    flush()
    return turns

def detect_roles(text):
    turns = parse_tadc_turns(text)
    roles = []
    seen = set()

    for role, _ in turns:
        key = role.strip().lower()
        if not key:
            continue
        if key not in seen:
            seen.add(key)
            roles.append(role.strip())

    return roles

def build_role_tokens(roles):
    tokens = [ROLE_PREFIX]
    for role in roles:
        tokens.append(f"{ROLE_PREFIX}{role}")
    return tokens

def normalize_tadc(text):
    """
    Rebuild transcript into a canonical format:
        <|role|>caine
        Hello!<|end|>
        <|role|>pomni
        Where am I?<|end|>
    """
    turns = parse_tadc_turns(text)
    parts = []
    for role, message in turns:
        parts.append(f"{ROLE_PREFIX}{role}\n{message}{END_TOKEN}\n")
    normalized = "".join(parts)
    print(f"   Normalized transcript: {len(normalized):,} characters")
    return normalized

def setup_tokenizer(tokenizer, text):
    print("Detecting TADC roles from transcript...")
    roles = detect_roles(text)
    role_tokens = build_role_tokens(roles)

    print(f"Found {len(roles)} unique roles:")
    for role in roles:
        print(f"   {ROLE_PREFIX}{role}")

    print(f"Adding {len(role_tokens)} TADC special tokens...")
    tokenizer.add_special_tokens({"additional_special_tokens": role_tokens})

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    tokenizer.model_max_length = MAX_LENGTH
    return tokenizer, roles, role_tokens

def create_long_examples(text, tokenizer):
    print("Tokenizing TADC transcript...")
    tokens = tokenizer.encode(text, add_special_tokens=False)
    print(f"   Total tokens: {len(tokens):,}")

    if len(tokens) <= MAX_LENGTH:
        chunks = [tokens]
    else:
        chunks = []
        step = MAX_LENGTH - CHUNK_OVERLAP
        for start in range(0, len(tokens), step):
            end = min(start + MAX_LENGTH, len(tokens))
            chunk = tokens[start:end]
            if chunk:
                chunks.append(chunk)
            if end >= len(tokens):
                break

    print(f"Created {len(chunks)} long training examples")
    if chunks:
        lengths = [len(x) for x in chunks]
        print(f"   Shortest: {min(lengths):,} tokens")
        print(f"   Longest:  {max(lengths):,} tokens")

    return Dataset.from_dict({"input_ids": chunks})

class TADCCollator:
    def __init__(self, tokenizer, max_length):
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __call__(self, features):
        input_ids = [f["input_ids"] for f in features]
        batch_size = len(input_ids)

        batch = torch.full(
            (batch_size, self.max_length),
            self.tokenizer.pad_token_id,
            dtype=torch.long,
        )
        attention_mask = torch.zeros(
            (batch_size, self.max_length),
            dtype=torch.long,
        )

        for i, ids in enumerate(input_ids):
            ids = ids[:self.max_length]
            length = len(ids)
            batch[i, :length] = torch.tensor(ids, dtype=torch.long)
            attention_mask[i, :length] = 1

        labels = batch.clone()
        labels[attention_mask == 0] = -100

        return {
            "input_ids": batch,
            "attention_mask": attention_mask,
            "labels": labels,
        }

def create_model(arch_choice, size, tokenizer):
    settings = SIZE_MODES[size]

    if arch_choice == "1":
        print(f"\nCreating GPT-2 TADC model ({size})")
        config = GPT2Config(
            vocab_size=len(tokenizer),
            n_positions=MAX_LENGTH,
            n_ctx=MAX_LENGTH,
            n_layer=settings["n_layer"],
            n_head=settings["n_head"],
            n_embd=settings["n_embd"],
            bos_token_id=tokenizer.bos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
        )
        model = GPT2LMHeadModel(config)

    elif arch_choice == "2":
        print(f"\nCreating LLaMA TADC model ({size})")
        config = LlamaConfig(
            vocab_size=len(tokenizer),
            max_position_embeddings=MAX_LENGTH,
            num_attention_heads=settings["n_head"],
            num_hidden_layers=settings["n_layer"],
            hidden_size=settings["hidden_size"],
            intermediate_size=settings["intermediate_size"],
            bos_token_id=tokenizer.bos_token_id,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.pad_token_id,
            attn_implementation="flash_attention_2",
        )
        model = LlamaForCausalLM(config)

    else:
        raise ValueError("Invalid architecture")

    model.resize_token_embeddings(len(tokenizer))
    model.config.use_cache = False
    model = model.to(device)
    return model

def train_tadc():
    text = load_tadc_transcript(DATASET_PATH)

    print("Architecture:")
    print("1) GPT-2")
    print("2) LLaMA")
    arch_choice = input("Select architecture: ").strip()

    print("Model size:")
    for name in SIZE_MODES:
        print(f"- {name}")
    size = input("Size: ").strip()

    if size not in SIZE_MODES:
        print("Invalid size")
        return

    resume_from_checkpoint = None
    if os.path.exists(OUTPUT_DIR) and os.listdir(OUTPUT_DIR):
        print("\nExisting model/checkpoint directory found.")
        resume_choice = input("Resume from checkpoint? (y/n): ").strip().lower()
        if resume_choice == "y":
            resume_from_checkpoint = input(
                "Enter checkpoint path (e.g. tadc/checkpoint-2000): "
            ).strip()
            print(f"Resuming from: {resume_from_checkpoint}")

    if arch_choice == "1":
        print("Loading GPT-2 tokenizer...")
        tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
    elif arch_choice == "2":
        print("Loading LLaMA tokenizer...")
        tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b")
    else:
        print("Invalid architecture")
        return

    tokenizer, roles, role_tokens = setup_tokenizer(tokenizer, text)

    print("Creating long TADC training examples...")
    normalized = normalize_tadc(text)
    dataset = create_long_examples(normalized, tokenizer)

    model = create_model(arch_choice, size, tokenizer)

    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

    print("=" * 60)
    print("MODEL PARAMETERS")
    print("=" * 60)
    print(f"Total:     {total_params:,}")
    print(f"Trainable: {trainable_params:,}")
    print(f"Approx:    {total_params / 1e6:.2f}M")

    collator = TADCCollator(tokenizer, MAX_LENGTH)

    use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported()
    use_fp16 = torch.cuda.is_available() and not use_bf16

    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=epoch_num,
        per_device_train_batch_size=6,
        gradient_accumulation_steps=4,
        gradient_checkpointing=True,
        bf16=use_bf16,
        fp16=use_fp16,
        learning_rate=3e-4,
        warmup_ratio=0.03,
        save_steps=50,
        save_total_limit=5,
        save_strategy="steps",
        logging_steps=5,
        report_to="none",
        dataloader_num_workers=0,
        remove_unused_columns=False,
        optim="adamw_torch",
        max_grad_norm=1.0,
    )

    wait_for_training_window()

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dataset,
        data_collator=collator,
        callbacks=[TimeWindowCallback()],
    )

    print("=" * 60)
    print("STARTING TADC TRAINING")
    print("=" * 60)
    print(f"Examples: {len(dataset)}")
    print(f"Context: {MAX_LENGTH:,} tokens")
    print(f"Overlap: {CHUNK_OVERLAP:,} tokens")
    print(f"Epochs: {epoch_num}")
    print(f"Roles: {len(roles)}")

    trainer.train(resume_from_checkpoint=resume_from_checkpoint)

    print("Saving final TADC model...")
    model.config.use_cache = True
    model.save_pretrained(OUTPUT_DIR, safe_serialization=True)
    tokenizer.save_pretrained(OUTPUT_DIR)

    with open(os.path.join(OUTPUT_DIR, "tadc_roles.json"), "w", encoding="utf-8") as f:
        json.dump(
            {"roles": roles, "special_tokens": role_tokens},
            f,
            indent=2,
            ensure_ascii=False,
        )

    print("TADC training complete!")
    print(f"Model saved to: {OUTPUT_DIR}")

def load_tadc_model(output_dir):
    print("Loading TADC model...")
    config_path = os.path.join(output_dir, "config.json")
    if not os.path.exists(config_path):
        raise FileNotFoundError(f"No config.json found in {output_dir}")

    with open(config_path, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    model_type = cfg.get("model_type", "")

    if model_type == "gpt2":
        tokenizer = GPT2TokenizerFast.from_pretrained(output_dir)
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.model_max_length = MAX_LENGTH
        if torch.cuda.is_available():
            model = GPT2LMHeadModel.from_pretrained(
                output_dir,
                torch_dtype=torch.float16,
                device_map="auto",
            )
        else:
            model = GPT2LMHeadModel.from_pretrained(output_dir)
            model = model.to(device)
    else:
        tokenizer = LlamaTokenizer.from_pretrained(output_dir)
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.model_max_length = MAX_LENGTH
        if torch.cuda.is_available():
            model = LlamaForCausalLM.from_pretrained(
                output_dir,
                torch_dtype=torch.float16,
                device_map="auto",
            )
        else:
            model = LlamaForCausalLM.from_pretrained(output_dir)
            model = model.to(device)

    model.eval()
    return tokenizer, model

def load_tadc_model(output_dir):
    print("Loading TADC model...")

    tokenizer = LlamaTokenizer.from_pretrained(output_dir)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.model_max_length = MAX_LENGTH

    roles_path = os.path.join(output_dir, "tadc_roles.json")
    roles = []

    if os.path.exists(roles_path):
        with open(roles_path, "r", encoding="utf-8") as f:
            meta = json.load(f)
        roles = meta.get("roles", [])

    if torch.cuda.is_available():
        model = LlamaForCausalLM.from_pretrained(
            output_dir,
            torch_dtype=torch.float16,
            device_map="auto",
        )
    else:
        model = LlamaForCausalLM.from_pretrained(output_dir)
        model = model.to(device)

    model.eval()
    return tokenizer, model, roles


def run_tadc():
    output_dir = input("Model path [tadc]: ").strip()
    if not output_dir:
        output_dir = "tadc"

    tokenizer, model, roles = load_tadc_model(output_dir)
    role_lookup = {r.lower(): r for r in roles}

    print("=" * 60)
    print("TADC MODEL RUNNER")
    print("=" * 60)
    print("Type 'exit' to quit.")
    print("Use the format:")
    print("Character: pomni")
    print("Message: Where am I?")

    while True:
        character_in = input("\nCharacter: ").strip()
        if character_in.lower() == "exit":
            break
        if not character_in:
            continue

        message = input("Message: ").strip()
        if message.lower() == "exit":
            break

        character = role_lookup.get(character_in.lower(), character_in)

        # Match the exact training format.
        prompt = f"<|role|><|{character}|>\n{message}\n"

        inputs = tokenizer(prompt, return_tensors="pt")
        model_device = next(model.parameters()).device
        inputs = {k: v.to(model_device) for k, v in inputs.items()}

        end_token_id = tokenizer.convert_tokens_to_ids(END_TOKEN)

        with torch.no_grad():
            output = model.generate(
                **inputs,
                max_new_tokens=120,
                do_sample=True,
                temperature=0.8,
                top_p=0.95,
                top_k=50,
                repetition_penalty=1.05,
                no_repeat_ngram_size=3,
                eos_token_id=end_token_id,
                pad_token_id=tokenizer.pad_token_id,
            )

        generated = output[0][inputs["input_ids"].shape[1]:]
        text = tokenizer.decode(generated, skip_special_tokens=False)

        if END_TOKEN in text:
            text = text.split(END_TOKEN, 1)[0]

        text = re.sub(r"<\|[^>]+?\|>", "", text).strip()

        print(text or "...")

def main():
    print("=" * 60)
    print("TADC TRAINER")
    print("=" * 60)

    while True:
        print("1) Train TADC model")
        print("2) Run TADC model")
        print("3) Exit")

        choice = input("Select: ").strip()

        if choice == "1":
            train_tadc()
        elif choice == "2":
            run_tadc()
        elif choice == "3":
            print("Bye")
            break
        else:
            print("Invalid choice")

if __name__ == "__main__":
    main()