modelmakerlibrary 0.1.0__py3-none-any.whl

src/model.py

@@ -0,0 +1,343 @@
+from .packages import torch, nn, json, pathlib, sft, plt, csv
+from .NeuralNetwork_create import TextGenerationModel, init_nn, init_text_generator
+from .dataset_function import load_dataset
+from accelerator import init_accelerator
+
+
+# Display the model specifications like parameters, batch size, learning rate, optimizer, loss function, and number of epochs.
+# Model params should be calculated based on the input size, output size, and hidden layers.
+def model_specs(parameters: dict, batch_size: int, lr_rate: float, optimizer: str, loss_function: str, epochs: int):
+    print("Model Specifications:")
+    input_size = parameters.get("input_size", 0)
+    output_size = parameters.get("output_size", 0)
+    hidden_layers = parameters.get("hidden_layers", 0)
+
+    if isinstance(hidden_layers, (list, tuple)):
+        hidden_sizes = list(hidden_layers)
+    elif isinstance(hidden_layers, int) and hidden_layers > 0:
+        hidden_sizes = [hidden_layers]
+    else:
+        hidden_sizes = []
+
+    layer_sizes = [input_size, *hidden_sizes, output_size]
+    parameter_count = sum(
+        (layer_sizes[index] + 1) * layer_sizes[index + 1]
+        for index in range(len(layer_sizes) - 1)
+        if isinstance(layer_sizes[index], int)
+        and isinstance(layer_sizes[index + 1], int)
+        and layer_sizes[index] >= 0
+        and layer_sizes[index + 1] >= 0
+    )
+
+    print(f"Input size: {input_size}")
+    print(f"Hidden layers: {hidden_layers}")
+    print(f"Output size: {output_size}")
+    print(f"Trainable parameters: {parameter_count}")
+    print(f"Batch size: {batch_size}")
+    print(f"Learning rate: {lr_rate}")
+    print(f"Optimizer: {optimizer}")
+    print(f"Loss function: {loss_function}")
+    print(f"Epochs: {epochs}")
+    
+
+# Save model in .pt ot .safetensors or .pkl format. The function takes the model, path, and format as input and saves the model accordingly.
+def model_save(model: nn.Module, path: str, format: str):
+    if format == "pt":
+        torch.save(model.state_dict(), path)
+        print(f"Model saved in .pt format at {path}")
+    elif format == "safetensors":
+        sft.torch.save_file(model.state_dict(), path)
+        print(f"Model saved in .safetensors format at {path}")
+    elif format == "pkl":
+        with open(path, 'wb') as f:
+            torch.save(model, f)
+        print(f"Model saved in .pkl format at {path}")
+    else:
+        raise ValueError("Unsupported format. Please choose from 'pt', 'safetensors', or 'pkl'.")
+
+
+
+# To quantize already trained models
+def model_quantization(model: nn.Module, model_file: str, out_path: str, quant_dtype: int):
+    if model is None:
+        raise ValueError("A model instance is required for quantization.")
+
+    if not quant_dtype:
+        raise ValueError("quant_dtype cannot be empty. Use values like 'int8', 'fp16', or 'bf16'.")
+
+    if model_file:
+        checkpoint = torch.load(model_file, map_location="cpu")
+
+        if isinstance(checkpoint, dict):
+            if "state_dict" in checkpoint:
+                state_dict = checkpoint["state_dict"]
+            elif all(isinstance(key, str) for key in checkpoint.keys()):
+                state_dict = checkpoint
+            else:
+                raise ValueError("Unsupported checkpoint format. Expected a state_dict or a dict containing 'state_dict'.")
+
+            model.load_state_dict(state_dict)
+        else:
+            raise TypeError("The model checkpoint must be a dictionary or a state dict.")
+
+    dtype_name = str(quant_dtype).strip().lower()
+
+    if dtype_name in ("", "int8", "qint8", "8"):
+        target_dtype = torch.qint8
+    elif dtype_name in ("float16", "fp16", "16"):
+        target_dtype = torch.float16
+    elif dtype_name in ("bfloat16", "bf16", "32"):
+        target_dtype = torch.bfloat16
+    elif isinstance(quant_dtype, int):
+        if quant_dtype == 8:
+            target_dtype = torch.qint8
+        elif quant_dtype == 16:
+            target_dtype = torch.float16
+        elif quant_dtype == 32:
+            target_dtype = torch.bfloat16
+        else:
+            raise ValueError("Unsupported quant_dtype. Use 'int8', 'fp16', or 'bf16'.")
+    else:
+        raise ValueError("Unsupported quant_dtype. Use 'int8', 'fp16', or 'bf16'.")
+
+    model.eval()
+
+    if target_dtype == torch.qint8:
+        quantized_model = torch.quantization.quantize_dynamic(
+            model,
+            {nn.Linear},
+            dtype=torch.qint8,
+        )
+        torch.save(quantized_model.state_dict(), out_path)
+        print(f"Model quantized to int8 and saved at {out_path}")
+        return quantized_model
+
+    model = model.to(dtype=target_dtype)
+    torch.save(model.state_dict(), out_path)
+    print(f"Model converted to {dtype_name.upper()} and saved at {out_path}")
+    return model
+
+
+
+def model_inference(model_type: str, path: str, test_data_path: str):
+    """Evaluate a saved model on labeled test data and display its accuracy."""
+    try:
+        checkpoint = torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        checkpoint = torch.load(path, map_location="cpu")
+
+    if isinstance(checkpoint, nn.Module):
+        model = checkpoint
+    else:
+        model = checkpoint.get("model") if isinstance(checkpoint, dict) else None
+        if not isinstance(model, nn.Module):
+            model = init_nn(model_type)
+            state_dict = checkpoint.get("state_dict", checkpoint) if isinstance(checkpoint, dict) else checkpoint
+            model.load_state_dict(state_dict)
+
+    suffix = pathlib.Path(test_data_path).suffix.lower()
+    if suffix in (".pt", ".pth"):
+        try:
+            test_data = torch.load(test_data_path, map_location="cpu", weights_only=False)
+        except TypeError:
+            test_data = torch.load(test_data_path, map_location="cpu")
+    elif suffix == ".json":
+        with open(test_data_path, "r", encoding="utf-8") as file:
+            test_data = json.load(file)
+    elif suffix == ".csv":
+        
+        with open(test_data_path, "r", newline="", encoding="utf-8") as file:
+            rows = list(csv.reader(file))
+        if rows:
+            try:
+                [float(value) for value in rows[0]]
+            except ValueError:
+                rows = rows[1:]
+        if not rows or any(len(row) < 2 for row in rows):
+            raise ValueError("CSV test data must contain feature columns and a label column.")
+        rows = [[float(value) for value in row] for row in rows]
+        test_data = ([row[:-1] for row in rows], [row[-1] for row in rows])
+    else:
+        raise ValueError("Test data must be a .pt, .pth, .json, or .csv file.")
+
+    if isinstance(test_data, dict):
+        features = next((test_data[key] for key in ("features", "inputs", "X", "data") if key in test_data), None)
+        labels = next((test_data[key] for key in ("labels", "targets", "y") if key in test_data), None)
+    elif isinstance(test_data, (list, tuple)) and len(test_data) == 2:
+        features, labels = test_data
+    elif isinstance(test_data, list) and test_data and isinstance(test_data[0], (list, tuple)):
+        features = [row[:-1] for row in test_data]
+        labels = [row[-1] for row in test_data]
+    else:
+        raise ValueError("Test data must contain both features and labels.")
+
+    if features is None or labels is None:
+        raise ValueError("Test data must include both features and labels.")
+    features = torch.as_tensor(features, dtype=torch.float32)
+    labels = torch.as_tensor(labels)
+    if features.ndim == 1:
+        features = features.unsqueeze(0)
+
+    model.eval()
+    with torch.no_grad():
+        outputs = model(features)
+        if isinstance(outputs, (tuple, list)):
+            outputs = outputs[0]
+        if outputs.ndim > 1 and outputs.shape[-1] > 1:
+            predictions = outputs.argmax(dim=1)
+        else:
+            predictions = (outputs.reshape(-1) >= 0.5).long()
+        accuracy = (predictions.reshape(-1) == labels.reshape(-1).long()).float().mean().item() * 100
+
+    print(f"Model accuracy: {accuracy:.2f}%")
+    
+    plt.figure()
+    plt.bar(["Accuracy"], [accuracy], color="steelblue")
+    plt.ylim(0, 100)
+    plt.ylabel("Performance (%)")
+    plt.title("Model Test Performance")
+    plt.text(0, accuracy, f"{accuracy:.2f}%", ha="center", va="bottom")
+    plt.tight_layout()
+    plt.show()
+    return accuracy
+
+
+def train_text_generator(
+    dataset_path: str,
+    output_path: str,
+    epochs: int = 10,
+    batch_size: int = 16,
+    steps_per_epoch: int = 10,
+    learning_rate: float = 3e-4,
+    device: str = "cpu",
+    max_context_length: int = 64,
+    max_new_tokens: int = 120,
+    prompt: str = "ModelMakerLibrary",
+) -> dict:
+    """Train and save a causal byte-level text model."""
+    for value, name in (
+        (epochs, "Epoch count"),
+        (batch_size, "Batch size"),
+        (steps_per_epoch, "Steps per epoch"),
+        (max_context_length, "Maximum context length"),
+        (max_new_tokens, "Maximum new tokens"),
+    ):
+        if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+            raise ValueError(f"{name} must be a positive integer.")
+    if isinstance(learning_rate, bool) or not isinstance(learning_rate, (int, float)) or learning_rate <= 0:
+        raise ValueError("Learning rate must be a positive number.")
+    if not isinstance(prompt, str) or not prompt:
+        raise ValueError("Prompt must be a non-empty string.")
+
+    dataset = load_dataset(dataset_path, model_type="text")
+    corpus = "\n".join(
+        text.strip()
+        for text, _ in dataset
+        if isinstance(text, str) and text.strip()
+    )
+    token_ids = torch.tensor(list(corpus.encode("utf-8")), dtype=torch.long)
+    if token_ids.numel() < 9:
+        raise ValueError("The training dataset must contain at least 9 UTF-8 bytes.")
+
+    context_length = min(max_context_length, max(8, token_ids.numel() // 3))
+    if token_ids.numel() <= context_length:
+        raise ValueError("The training dataset must be longer than the model context.")
+    target_device = init_accelerator(device)
+    model = init_text_generator(context_length=context_length).to(target_device)
+    parameter_count = sum(parameter.numel() for parameter in model.parameters())
+    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
+    losses = []
+
+    model.train()
+    for epoch in range(1, epochs + 1):
+        total_loss = 0.0
+        for _ in range(steps_per_epoch):
+            starts = torch.randint(token_ids.numel() - context_length, (batch_size,))
+            offsets = torch.arange(context_length)
+            indices = starts.unsqueeze(1) + offsets
+            inputs = token_ids[indices].to(target_device)
+            targets = token_ids[indices + 1].to(target_device)
+
+            optimizer.zero_grad(set_to_none=True)
+            logits = model(inputs)
+            loss = nn.functional.cross_entropy(
+                logits.reshape(-1, model.vocab_size),
+                targets.reshape(-1),
+            )
+            loss.backward()
+            optimizer.step()
+            total_loss += loss.item()
+        average_loss = total_loss / steps_per_epoch
+        losses.append(average_loss)
+        print(f"Epoch {epoch:02d}/{epochs}: loss={average_loss:.4f}")
+
+    checkpoint_path = pathlib.Path(output_path).expanduser().resolve()
+    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
+    torch.save(
+        {
+            "state_dict": model.state_dict(),
+            "context_length": model.context_length,
+            "vocab_size": model.vocab_size,
+            "model_width": model.model_width,
+            "attention_heads": model.attention_heads,
+            "transformer_layers": model.transformer_layers,
+            "parameter_count": parameter_count,
+            "losses": losses,
+        },
+        checkpoint_path,
+    )
+    sample = generate_text(model, prompt, max_new_tokens=max_new_tokens)
+    return {
+        "checkpoint_path": str(checkpoint_path),
+        "parameter_count": parameter_count,
+        "context_length": context_length,
+        "losses": losses,
+        "sample": sample,
+    }
+
+
+def load_text_generator(path: str, device: str = "cpu") -> TextGenerationModel:
+    """Load a text-generation checkpoint saved by train_text_generator."""
+    target_device = init_accelerator(device)
+    try:
+        checkpoint = torch.load(path, map_location=target_device, weights_only=True)
+    except TypeError:
+        checkpoint = torch.load(path, map_location=target_device)
+    if not isinstance(checkpoint, dict) or "state_dict" not in checkpoint:
+        raise ValueError("Unsupported text-generator checkpoint.")
+    model = init_text_generator(
+        context_length=checkpoint["context_length"],
+        vocab_size=checkpoint["vocab_size"],
+        model_width=checkpoint["model_width"],
+        attention_heads=checkpoint["attention_heads"],
+        transformer_layers=checkpoint["transformer_layers"],
+    )
+    model.load_state_dict(checkpoint["state_dict"])
+    model.to(target_device)
+    model.eval()
+    return model
+
+
+def generate_text(
+    model: TextGenerationModel,
+    prompt: str,
+    max_new_tokens: int = 120,
+) -> str:
+    """Generate text from a prompt using a library-created language model."""
+    if not isinstance(model, TextGenerationModel):
+        raise TypeError("Model must be a TextGenerationModel created by ModelMakerLibrary.")
+    if not isinstance(prompt, str) or not prompt:
+        raise ValueError("Prompt must be a non-empty string.")
+    if isinstance(max_new_tokens, bool) or not isinstance(max_new_tokens, int) or max_new_tokens <= 0:
+        raise ValueError("Maximum new tokens must be a positive integer.")
+
+    device = next(model.parameters()).device
+    generated = torch.tensor(list(prompt.encode("utf-8")), dtype=torch.long, device=device)
+    model.eval()
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            context = generated[-model.context_length :].unsqueeze(0)
+            logits = model(context)[0, -1]
+            next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
+            generated = torch.cat((generated, next_token))
+    return bytes(generated.tolist()).decode("utf-8", errors="replace")

src/packages.py

@@ -0,0 +1,17 @@
+import torch
+from torch import optim
+from torch import nn
+from torch.utils.data import DataLoader
+from torchvision import datasets
+from torchvision.transforms import v2
+import json 
+import csv
+import tsv
+import pandas as pd
+import transformers
+import os
+import tokenizers
+import pathlib
+from accelerator import *
+import matplotlib.pyplot as plt
+import safetensors.torch as sft

src/tokenizers_function.py

@@ -0,0 +1,109 @@
+from src.packages import os, transformers, torch
+from src.dataset_function import init_dataset
+
+def init_tokenizer(tokenizer_name, tokenizer_path=None, token=False):
+    """Load a tokenizer, ignoring saved Hugging Face credentials unless provided."""
+    if tokenizer_name is None:
+        raise ValueError("Tokenizer name must not be None.")
+    if tokenizer_name == "bert":
+        if tokenizer_path is not None:
+            tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path, token=token)
+        else:
+            tokenizer = transformers.AutoTokenizer.from_pretrained('bert-base-uncased', token=token)
+    elif tokenizer_name == "gpt2":
+        if tokenizer_path is not None:
+            tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path, token=token)
+        else:
+            tokenizer = transformers.AutoTokenizer.from_pretrained('gpt2', token=token)
+    elif tokenizer_name == "roberta":
+        if tokenizer_path is not None:
+            tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path, token=token)
+        else:
+            tokenizer = transformers.AutoTokenizer.from_pretrained('roberta-base', token=token)
+    elif tokenizer_name == "custom":
+        if tokenizer_path is None:
+            raise ValueError("Tokenizer path must be provided for custom tokenizer.")
+        tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_path, token=token)
+    else:
+        raise ValueError("Unsupported tokenizer name. Supported names are: 'bert', 'gpt2', 'roberta', 'custom'.")
+    return tokenizer
+
+
+def tokenize_dataset(tokenizer, dataset_path, output_path, max_length=512):
+    if tokenizer is None:
+        raise ValueError("Tokenizer must not be None.")
+    if dataset_path is None or output_path is None:
+        raise ValueError("Dataset path and output path must not be None.")
+    if isinstance(max_length, bool) or not isinstance(max_length, int) or max_length <= 0:
+        raise ValueError("Maximum length must be a positive integer.")
+    if not os.path.exists(dataset_path):
+        raise ValueError("Dataset path does not exist.")
+
+    output_path = os.fspath(output_path)
+    absolute_output_path = os.path.abspath(output_path)
+    os.makedirs(os.path.dirname(absolute_output_path), exist_ok=True)
+    torch.save([], output_path)
+    print(f"Created token output file at '{absolute_output_path}'.")
+
+    _, dataset_path = init_dataset("text", os.fspath(dataset_path))
+    text_files = []
+    found_files = []
+    supported_extensions = {".txt", ".json", ".csv", ".tsv"}
+    if os.path.isfile(dataset_path):
+        found_files.append(dataset_path)
+        if os.path.splitext(dataset_path)[1].lower() in supported_extensions:
+            text_files.append(dataset_path)
+    else:
+        for root, _, files in os.walk(dataset_path):
+            for filename in files:
+                file_path = os.path.join(root, filename)
+                found_files.append(file_path)
+                if os.path.splitext(filename)[1].lower() in supported_extensions:
+                    text_files.append(file_path)
+    if not text_files:
+        details = "; ".join(
+            f"{os.path.relpath(path, dataset_path) if os.path.isdir(dataset_path) else os.path.basename(path)} "
+            f"({os.path.getsize(path)} bytes)"
+            for path in found_files[:5]
+        )
+        if len(found_files) > 5:
+            details += f"; and {len(found_files) - 5} more file(s)"
+        found_message = f" Files found: {details}." if details else " The directory contains no files."
+        raise ValueError(
+            f"Cannot tokenize dataset at '{os.path.abspath(dataset_path)}': no supported text data "
+            "was found. Supported inputs are UTF-8 .txt, JSON, CSV, and TSV files."
+            f"{found_message} Check that the input path points to the source text dataset or a "
+            "text dataset exported with dataset_convert."
+        )
+    tokenized_data = []
+
+    for text_file in sorted(text_files):
+        extension = os.path.splitext(text_file)[1].lower()
+        with open(text_file, encoding="utf-8", newline="") as file:
+            if extension == ".txt":
+                texts = [file.read()]
+            elif extension == ".json":
+                records = json.load(file)
+                if isinstance(records, dict):
+                    records = records.get("data", [records])
+                if not isinstance(records, list):
+                    raise ValueError(f"Expected a JSON object or list of records in '{text_file}'.")
+                texts = [
+                    record["text"] for record in records
+                    if isinstance(record, dict) and isinstance(record.get("text"), str)
+                ]
+            else:
+                reader = csv.DictReader(file, delimiter="\t" if extension == ".tsv" else ",")
+                if not reader.fieldnames or "text" not in reader.fieldnames:
+                    raise ValueError(f"Expected a 'text' column in '{text_file}'.")
+                texts = [row["text"] for row in reader if row.get("text")]
+        if not texts:
+            raise ValueError(f"No text records found in '{text_file}'.")
+        for text in texts:
+            tokenized_batch = tokenizer([text], padding=True, truncation=True, max_length=max_length, return_tensors='pt')
+            tokenized_data.append(tokenized_batch)
+
+    # Save the tokenized data to the output path
+    torch.save(tokenized_data, output_path)
+    print(f"Tokenized {len(tokenized_data)} text record(s) and saved them to '{absolute_output_path}'.")
+    return output_path

src/train.py

@@ -0,0 +1,36 @@
+"""Example training run using the ModelMakerLibrary public API."""
+
+from pathlib import Path
+
+import modelMakerLibrary as mml
+
+
+def main():
+    project_dir = Path(__file__).resolve().parents[1]
+    result = mml.train_text_generator(
+        dataset_path=str(project_dir / "data" / "train"),
+        output_path=str(project_dir / "data" / "text_generator.pt"),
+        epochs=10,
+        batch_size=16,
+        steps_per_epoch=10,
+        learning_rate=3e-4,
+        device="cpu",
+        max_context_length=64,
+        max_new_tokens=120,
+    )
+    parameter_count = result["parameter_count"]
+    print(
+        f"Trainable parameters: {parameter_count:,} "
+        f"({parameter_count / 1_000_000:.3f}M)"
+    )
+    if not 90_000 <= parameter_count <= 110_000:
+        raise RuntimeError(
+            f"Expected about 0.1M parameters, got {parameter_count:,}."
+        )
+    print(f"Saved checkpoint: {result['checkpoint_path']}")
+    print("Sample generation:")
+    print(result["sample"].encode("ascii", errors="backslashreplace").decode("ascii"))
+
+
+if __name__ == "__main__":
+    main()