scout-ai 0.2.0 → 1.0.0

data/python/scout_ai/huggingface/eval.py

@@ -0,0 +1,60 @@
+def forward(model, features):
+    return model(**features)
+
+def get_logits(predictions):
+    logits = predictions["logits"]
+    return [v.detach().cpu().numpy() for v in logits]
+
+def eval_model(model, tokenizer, texts, return_logits=True):
+    features = tokenizer(texts, return_tensors='pt', truncation=True).to(model.device)
+    model.eval()
+    predictions = forward(model, features)
+    if return_logits:
+        return get_logits(predictions)
+    return predictions
+
+def eval_causal_lm_chat(
+    model, tokenizer, messages,
+    chat_template=None,
+    chat_template_kwargs=None,
+    generation_kwargs=None
+):
+    """
+    Evaluate a CausalLM model given chat messages. Uses tokenizer's chat template by default.
+
+    Args:
+        model: Huggingface CausalLM
+        tokenizer: Huggingface tokenizer
+        messages: List[Dict[str, str]] (OpenAI API style, 'role' and 'content')
+        chat_template: (Optional) Override string for the chat template.
+        chat_template_kwargs: (Optional) Dict, kwargs for apply_chat_template (like tokenize, add_generation_prompt, etc).
+        generation_kwargs: (Optional) Dict for model.generate
+
+    Returns:
+        Generated text (or list, depending on settings).
+    """
+    chat_template_kwargs = chat_template_kwargs or {}
+    generation_kwargs = generation_kwargs or {}
+
+    # If the tokenizer has a chat template (HF 4.34+)
+    if hasattr(tokenizer, "___apply_chat_template"):
+        kwargs = dict(add_generation_prompt=True, tokenize=False)
+        kwargs.update(chat_template_kwargs)
+        if chat_template is not None:
+            # Override the template (may require tokenizer._chat_template)
+            tokenizer._chat_template = chat_template
+        prompt = tokenizer.apply_chat_template(messages, **kwargs)
+    else:
+        # Fallback: simple concatenation
+        prompt = "\n".join([msg['content'] for msg in messages])
+
+    # Tokenize as usual
+    inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
+    model.eval()
+    # Use generate
+    output_ids = model.generate(**inputs, **generation_kwargs)
+    # Decode only the newly generated tokens (not the prompt)
+    output_text = tokenizer.decode(
+        output_ids[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True
+    )
+    return output_text

data/python/scout_ai/huggingface/model.py

@@ -0,0 +1,29 @@
+# huggingface_model.py
+import importlib
+from typing import Optional, Any
+
+def import_module_class(module: str, class_name: str) -> Any:
+    """Dynamically import a class from a module."""
+    mod = importlib.import_module(module)
+    return getattr(mod, class_name)
+
+def load_model(task: Optional[str], checkpoint: str, **kwargs) -> Any:
+    """Load a Huggingface model by task and checkpoint"""
+    if task is None or task.lower() == 'embedding':
+        model_class = import_module_class('transformers', 'AutoModel')
+    elif ":" in task:
+        module, class_name = task.split(":")
+        model_class = import_module_class(module, class_name)
+    else:
+        model_class = import_module_class('transformers', f'AutoModelFor{task}')
+    return model_class.from_pretrained(checkpoint, **kwargs)
+
+def load_tokenizer(checkpoint: str, **kwargs) -> Any:
+    """Load a Huggingface tokenizer"""
+    tokenizer_class = import_module_class('transformers', 'AutoTokenizer')
+    return tokenizer_class.from_pretrained(checkpoint, **kwargs)
+
+def load_model_and_tokenizer(task: Optional[str], checkpoint: str, **kwargs):
+    model = load_model(task, checkpoint, **kwargs)
+    tokenizer = load_tokenizer(checkpoint, **kwargs)
+    return model, tokenizer

data/python/scout_ai/huggingface/rlhf.py

@@ -0,0 +1,83 @@
+from trl import PPOTrainer, AutoModelForCausalLMWithValueHead, PPOConfig
+import torch
+import scout_ai
+
+from copy import deepcopy
+from datasets import Dataset
+
+
+class PPOTrainerWithPrecomputedReward(PPOTrainer):
+    def get_rewards(self, **kwargs):
+        return torch.tensor(self.train_dataset['reward'], dtype=torch.float32)
+
+def train_rlhf(path, tokenizer, pairs, rewards, config=None, generation_config=None):
+    """
+    pairs: List of tuples (messages, response)
+        - messages: List[Dict[str, str]] (OpenAI/chatML-style messages)
+        - response: string (the model output to be rewarded)
+    """
+    config = config or {}
+    device = scout_ai.device()
+    device = 'cuda'
+
+    tokenizer.padding_side = "left"
+    tokenizer.pad_token = tokenizer.eos_token
+
+    prompts, responses = [], []
+    for pair in pairs:
+        messages, response = pair
+        # Ensure tokenizer supports chat template (HF >=4.34)
+        if hasattr(tokenizer, 'apply_chat_template'):
+            # Use default: add_generation_prompt needed for LLMs like Llama, Mistral, etc
+            prompt = tokenizer.apply_chat_template(
+                messages, add_generation_prompt=True, tokenize=False
+            )
+        else:
+            # Fallback: join user/assistant messages
+            prompt = "\n".join(msg['content'] for msg in messages)
+        prompts.append(prompt)
+        responses.append(response)
+
+    train_dataset = Dataset.from_dict({'prompt': prompts, 'response': responses, 'reward': rewards})
+
+    # Wrap model with Value Head for PPO
+    from trl import PPOTrainer, AutoModelForCausalLMWithValueHead, PPOConfig
+    model = AutoModelForCausalLMWithValueHead.from_pretrained(path)
+    model.to(device)
+
+    from transformers import GenerationConfig
+    
+    generation_config = GenerationConfig()
+
+    ppo_config = PPOConfig(
+        batch_size=config.get('batch_size', 4),
+        learning_rate=config.get('learning_rate', 1e-5),
+        mini_batch_size=config.get('mini_batch_size', 1),
+        gradient_accumulation_steps=1,
+    )
+
+    model.base_model_prefix = 'model'
+
+    ref_model = deepcopy(model)
+    ref_model.to(device)
+
+    model.generation_config=generation_config
+
+    print(model)
+    print(ref_model)
+
+    ppo_trainer = PPOTrainerWithPrecomputedReward(
+        args=ppo_config,
+        model=model,
+        ref_model=ref_model,
+        reward_model=model,  # dummy
+        value_model=model,  # dummy
+        train_dataset=train_dataset,
+        processing_class=None,
+    )
+
+    
+    print("Step")
+    stats = ppo_trainer.train(prompts, responses, rewards)
+    model.save
+    return stats

data/python/scout_ai/huggingface/train/__init__.py

@@ -0,0 +1,34 @@
+from transformers import TrainingArguments, Trainer
+from typing import Any
+from ..data import json_dataset, tsv_dataset, tokenize_dataset
+
+def training_args(*args, **kwargs) -> TrainingArguments:
+    return TrainingArguments(*args, **kwargs)
+
+def train_model(model: Any, tokenizer: Any, training_args: TrainingArguments, dataset: Any, class_weights=None, **kwargs):
+    for param in model.parameters():
+        param.data = param.data.contiguous()
+
+    if (isinstance(dataset, str)):
+        if (dataset.endswith('.json')):
+            tokenized_dataset = json_dataset(tokenizer, dataset)
+        else:
+            tokenized_dataset = tsv_dataset(tokenizer, dataset)
+    else:
+        tokenized_dataset = tokenize_dataset(tokenizer, dataset)
+
+    if class_weights is not None:
+        import torch
+        from torch import nn
+        class WeightTrainer(Trainer):
+            def compute_loss(self, model, inputs, return_outputs=False):
+                labels = inputs.get("labels")
+                outputs = model(**inputs)
+                logits = outputs.get('logits')
+                loss_fct = nn.CrossEntropyLoss(weight=torch.tensor(class_weights).to(model.device))
+                loss = loss_fct(logits.view(-1, model.config.num_labels), labels.view(-1))
+                return (loss, outputs) if return_outputs else loss
+        trainer = WeightTrainer(model, training_args, train_dataset=tokenized_dataset["train"], tokenizer=tokenizer, **kwargs)
+    else:
+        trainer = Trainer(model, training_args, train_dataset=tokenized_dataset["train"], tokenizer=tokenizer, **kwargs)
+    trainer.train()

data/python/scout_ai/huggingface/train/next_token.py

@@ -0,0 +1,315 @@
+import os
+import math
+import time
+import shutil
+import random
+from dataclasses import dataclass
+from typing import List, Optional, Dict, Any, Union
+
+import torch
+from torch.utils.data import DataLoader
+from datasets import Dataset, load_dataset
+
+from transformers import (
+    PreTrainedModel, 
+    PreTrainedTokenizer, 
+    get_scheduler, 
+    DataCollatorForLanguageModeling
+)
+from torch.optim import AdamW
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+def set_seed(seed: int):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    try:
+        import numpy as np
+        np.random.seed(seed)
+    except ImportError:
+        pass
+
+@dataclass
+class TrainingState:
+    global_step: int = 0
+    best_eval_loss: float = float("inf")
+
+def tokenize_function(examples, tokenizer, max_seq_length):
+    # examples: dict with key 'text' or single texts
+    # Always output input_ids and attention_mask
+    output = tokenizer(
+        examples["text"] if "text" in examples else examples,
+        truncation=True,
+        padding="max_length",
+        max_length=max_seq_length,
+        return_attention_mask=True,
+    )
+    output["labels"] = output["input_ids"].copy()
+    return output
+
+def group_texts(examples, block_size):
+    # For paragraph-based datasets: simply return; for huge files, use this.
+    concatenated = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated[list(examples.keys())[0]])
+    # Drop the small remainder
+    total_length = (total_length // block_size) * block_size
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated.items()
+    }
+    return result
+
+def train_next_token(
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    dataset: Union[List[str], Dataset],
+    *,
+    output_dir: str,
+    eval_dataset: Optional[Union[List[str], Dataset]] = None,
+    max_seq_length: int = 2048,
+    batch_size: int = 8,
+    gradient_accumulation_steps: int = 1,
+    num_train_epochs: int = 3,
+    learning_rate: float = 1e-4,
+    weight_decay: float = 0.01,
+    lr_scheduler_type: str = "linear",
+    warmup_steps: int = 0,
+    logging_steps: int = 50,
+    eval_steps: int = 200,
+    save_steps: int = 500,
+    save_total_limit: int = 3,
+    fp16: bool = False,
+    bf16: bool = False,
+    max_train_steps: int = None,
+    seed: int = 42,
+    report_to: str = "none",  # or "wandb", "tensorboard"
+    use_lora: bool = False,
+    lora_config: Optional[dict] = None,
+    resume_from_checkpoint: str = None,
+    callbacks: Optional[List] = None,
+    device_map: str = "auto",
+    dataloader_num_workers: int = 4,
+    group_by_length: bool = False,
+    description: str = "",
+):
+    """
+    Fine-tunes a causal LM for next-token prediction.
+    """
+    #assert isinstance(model, PreTrainedModel), "Model must be a HuggingFace PreTrainedModel"
+    #assert isinstance(tokenizer, PreTrainedTokenizer), "Tokenizer must be a HuggingFace PreTrainedTokenizer"
+    assert isinstance(dataset, (list, Dataset)), "Dataset must be a HuggingFace Dataset or a list of texts"
+
+    set_seed(seed)
+    os.makedirs(output_dir, exist_ok=True)
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    n_gpus = torch.cuda.device_count()
+
+    if resume_from_checkpoint:
+        logger.info(f"Loading checkpoint from {resume_from_checkpoint}")
+        model.load_state_dict(torch.load(os.path.join(resume_from_checkpoint, "pytorch_model.bin")))
+
+    model.to(device)
+
+    if fp16:
+        scaler = torch.cuda.amp.GradScaler()
+    else:
+        scaler = None
+
+    # 1. Prepare Dataset
+    if isinstance(dataset, list):
+        dataset = Dataset.from_dict({"text": dataset})
+
+    if eval_dataset is not None and isinstance(eval_dataset, list):
+        eval_dataset = Dataset.from_dict({"text": eval_dataset})
+
+    # Tokenization and formatting
+    def preprocess(examples):
+        return tokenize_function(examples, tokenizer, max_seq_length)
+
+    dataset = dataset.map(preprocess, batched=True, remove_columns=list(dataset.column_names))
+    if eval_dataset is not None:
+        eval_dataset = eval_dataset.map(preprocess, batched=True, remove_columns=list(eval_dataset.column_names))
+
+    # 2. Loader & Collator
+    data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+
+    train_loader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=True,
+        collate_fn=data_collator,
+        num_workers=dataloader_num_workers,
+        drop_last=True,
+    )
+    eval_loader = None
+    if eval_dataset is not None:
+        eval_loader = DataLoader(
+            eval_dataset, 
+            batch_size=batch_size,
+            shuffle=False,
+            collate_fn=data_collator, 
+            num_workers=dataloader_num_workers,
+        )
+
+    # 3. Optimizer & Scheduler
+    no_decay = ["bias", "LayerNorm.weight"]
+    grouped_params = [
+        {
+            "params": [
+                p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": weight_decay,
+        },
+        {
+            "params": [
+                p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)
+            ],
+            "weight_decay": 0.0,
+        },
+    ]
+
+    optimizer = AdamW(grouped_params, lr=learning_rate)
+
+    total_train_steps = (
+        max_train_steps if max_train_steps is not None
+        else (len(train_loader) * num_train_epochs) // gradient_accumulation_steps
+    )
+
+    lr_scheduler = get_scheduler(
+        lr_scheduler_type,
+        optimizer=optimizer,
+        num_warmup_steps=warmup_steps,
+        num_training_steps=total_train_steps,
+    )
+
+    # 4. LoRA/PEFT Support (placeholder)
+    if use_lora:
+        logger.warning("PEFT/LoRA integration not yet implemented. Skipping.")
+
+    # 5. Checkpoint Management
+    saved_checkpoints = []
+
+    # 6. Training Loop
+    state = TrainingState()
+    model.train()
+    start_time = time.time()
+    for epoch in range(num_train_epochs):
+        logger.info(f"Epoch {epoch+1}/{num_train_epochs}")
+        for step, batch in enumerate(train_loader):
+            true_step = state.global_step + 1
+            batch = {k: v.to(device) for k, v in batch.items()}
+            with torch.cuda.amp.autocast(dtype=torch.float16 if fp16 else torch.bfloat16 if bf16 else torch.float32, enabled=(fp16 or bf16)):
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+
+            if fp16:
+                scaler.scale(loss).backward()
+            else:
+                loss.backward()
+
+            if true_step % gradient_accumulation_steps == 0:
+                if fp16:
+                    scaler.step(optimizer)
+                    scaler.update()
+                else:
+                    optimizer.step()
+                optimizer.zero_grad()
+                lr_scheduler.step()
+
+            if true_step % logging_steps == 0:
+                logger.info(f"Step {true_step}: loss {loss.item() * gradient_accumulation_steps:.4f}")
+
+            if eval_loader is not None and true_step % eval_steps == 0:
+                eval_loss = evaluate(model, eval_loader, device, fp16, bf16)
+                logger.info(f"Step {true_step}: eval_loss {eval_loss:.4f}, ppl {math.exp(eval_loss):.2f}")
+                # Save best
+                if eval_loss < state.best_eval_loss:
+                    state.best_eval_loss = eval_loss
+                    save_checkpoint(model, output_dir, f"best")
+            if true_step % save_steps == 0:
+                ckpt_dir = save_checkpoint(model, output_dir, f"step-{true_step}")
+                saved_checkpoints.append(ckpt_dir)
+                # Cleanup
+                if len(saved_checkpoints) > save_total_limit:
+                    old = saved_checkpoints.pop(0)
+                    shutil.rmtree(old, ignore_errors=True)
+            state.global_step = true_step
+            if max_train_steps is not None and true_step >= max_train_steps:
+                break
+        # End-of-epoch eval/save
+        if eval_loader is not None:
+            eval_loss = evaluate(model, eval_loader, device, fp16, bf16)
+            logger.info(f"Epoch {epoch+1} end: eval_loss {eval_loss:.4f}, ppl {math.exp(eval_loss):.2f}")
+            if eval_loss < state.best_eval_loss:
+                state.best_eval_loss = eval_loss
+                save_checkpoint(model, output_dir, "best")
+        save_checkpoint(model, output_dir, f"epoch-{epoch+1}")
+    logger.info(f"Training completed in {time.time() - start_time:.2f} sec on {device}")
+
+def evaluate(model, eval_loader, device, fp16, bf16):
+    model.eval()
+    losses = []
+    for batch in eval_loader:
+        batch = {k: v.to(device) for k, v in batch.items()}
+        with torch.no_grad():
+            with torch.cuda.amp.autocast(dtype=torch.float16 if fp16 else torch.bfloat16 if bf16 else torch.float32, enabled=(fp16 or bf16)):
+                outputs = model(**batch)
+            losses.append(outputs.loss.item())
+    model.train()
+    return sum(losses) / len(losses)
+
+def save_checkpoint(model, output_dir, tag):
+    output_ckpt_dir = os.path.join(output_dir, tag)
+    os.makedirs(output_ckpt_dir, exist_ok=True)
+    model.save_pretrained(output_ckpt_dir)
+    return output_ckpt_dir
+
+def main():
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+
+    # Example tiny dataset: few sentences
+    train_texts = [
+        "The quick brown fox jumps over the lazy dog.",
+        "Artificial intelligence is the future.",
+        "Llama models are great for language tasks.",
+        "Open source is important for research.",
+    ]
+    eval_texts = [
+        "Transformers enable powerful NLP models.",
+        "Fine-tuning improves performance."
+    ]
+
+    #model_name = "openlm-research/open_llama_3b"  # Replace with your local/other HF Llama checkpoint as needed
+    model_name = "distilgpt2"  # Replace with your local/other HF Llama checkpoint as needed
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+    # Make sure tokenizer pads on right for causal LMs (Llama does not have pad by default)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    train_next_token(
+        model=model,
+        tokenizer=tokenizer,
+        dataset=train_texts,
+        output_dir="./output_test",
+        eval_dataset=eval_texts,
+        max_seq_length=32,
+        batch_size=2,
+        num_train_epochs=1,
+        gradient_accumulation_steps=1,
+        learning_rate=5e-5,
+        fp16=False,  # Change to True if running on GPU with enough VRAM
+        bf16=False,
+        logging_steps=1,
+        eval_steps=2,
+        save_steps=10
+    )
+
+if __name__ == "__main__":
+    main()

data/python/scout_ai/language_model.py

@@ -0,0 +1,70 @@
+def group_texts(examples):
+    # Concatenate all texts.
+    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
+    total_length = len(concatenated_examples[list(examples.keys())[0]])
+    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
+        # customize this part to your needs.
+    total_length = (total_length // block_size) * block_size
+    # Split by chunks of max_len.
+    result = {
+        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
+        for k, t in concatenated_examples.items()
+    }
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+def whole_word_masking_data_collator(features):
+    from transformers import default_data_collator
+    for feature in features:
+        word_ids = feature.pop("word_ids")
+
+        # Create a map between words and corresponding token indices
+        mapping = collections.defaultdict(list)
+        current_word_index = -1
+        current_word = None
+        for idx, word_id in enumerate(word_ids):
+            if word_id is not None:
+                if word_id != current_word:
+                    current_word = word_id
+                    current_word_index += 1
+                mapping[current_word_index].append(idx)
+
+        # Randomly mask words
+        mask = np.random.binomial(1, wwm_probability, (len(mapping),))
+        input_ids = feature["input_ids"]
+        labels = feature["labels"]
+        new_labels = [-100] * len(labels)
+        for word_id in np.where(mask)[0]:
+            word_id = word_id.item()
+            for idx in mapping[word_id]:
+                new_labels[idx] = labels[idx]
+                input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
+
+    return default_data_collator(features)
+
+if __name__ == "__main__2":
+
+    from transformers import AutoModelForMaskedLM
+    from transformers import AutoTokenizer
+    import torch
+
+    model_checkpoint = "distilbert-base-uncased"
+    model = AutoModelForMaskedLM.from_pretrained(model_checkpoint)
+    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+
+    text = "This is a great [MASK]."
+
+    inputs = tokenizer(text, return_tensors="pt")
+    token_logits = model(**inputs).logits
+    # Find the location of [MASK] and extract its logits
+    mask_token_index = torch.where(inputs["input_ids"] == tokenizer.mask_token_id)[1]
+    mask_token_logits = token_logits[0, mask_token_index, :]
+    # Pick the [MASK] candidates with the highest logits
+    top_5_tokens = torch.topk(mask_token_logits, 5, dim=1).indices[0].tolist()
+
+    for token in top_5_tokens:
+        print(f"'>>> {text.replace(tokenizer.mask_token, tokenizer.decode([token]))}'")
+
+
+

data/python/scout_ai/util.py

@@ -0,0 +1,32 @@
+import random
+import torch
+import numpy
+
+def set_seed(seed):
+    """
+    Set seed in several backends
+    """
+    random.seed(seed)
+    numpy.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+def deterministic():
+    """
+    Ensure that all operations are deterministic on GPU (if used) for
+    reproducibility
+    """
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def device():
+    return torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
+
+def data_directory():
+    from pathlib import Path
+    print(Path.home())
+
+def model_device(model):
+    return next(model.parameters()).device