Re-adiciona model.safetensors via LFS

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,96 @@
+# Byte‑compiled / build artifacts
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+src/checkpoints/
+mlflow/
+postgres-temp/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+*.pth
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+
+# IDEs and editors
+.idea/
+*.iml
+.vscode/
+*.sublime-project
+*.sublime-workspace
+
+# Docker
+docker-compose.override.yml
+.docker/
+
+# Environment / virtualenv
+.env
+.venv/
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Poetry
+poetry.lock
+poetry.toml
+
+# Pyproject (se configurar como público; se privado, remova)
+# pyproject.toml
+
+# Lock files
+Pipfile.lock
+
+# Runtime data
+*.pid
+*.seed
+*.log
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# VS Code settings
+.vscode/
+
+# OS files
+.DS_Store
+Thumbs.db

Dockerfile

@@ -0,0 +1,3 @@
+FROM ghcr.io/mlflow/mlflow:v2.21.3
+
+RUN pip install boto3

README.md

@@ -0,0 +1,7 @@
+---
+library_name: transformers
+license: apache-2.0
+pipeline_tag: text-generation
+base_model:
+- bobboyms/tynerox
+---

config.json

@@ -0,0 +1,19 @@
+{
+  "architectures": [
+    "TyneRoxModel"
+  ],
+  "causal": true,
+  "d_model": 1024,
+  "dropout": 0.1,
+  "layer_norm_eps": 1e-05,
+  "max_position_embeddings": 2048,
+  "model_type": "tynerox",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.51.3",
+  "vocab_size": 36010,
+  "window_size": 512
+}

docker-compose.yml

@@ -0,0 +1,45 @@
+version: '3'
+services:
+  mlflow:
+#    image: ghcr.io/mlflow/mlflow:v2.21.3
+#    image: ghcr.io/mlflow/mlflow:v3.0.0rc0
+    build: .
+    ports:
+      - "5000:5000"
+    environment:
+      - MLFLOW_ARTIFACT_ROOT=s3://1hh-mlflow/artifacts
+      - MLFLOW_TRACKING_URI=http://mlflow:5000
+    volumes:
+      - ./mlflow:/mlflow
+    command: mlflow server --host 0.0.0.0 --port 5000 --backend-store-uri sqlite:///mlflow.db --default-artifact-root s3://1hh-mlflow/artifacts
+# s3://1hh-mlflow/artifacts/
+  postgres:
+    image: postgres:14
+    environment:
+      - POSTGRES_USER=mlflow
+      - POSTGRES_PASSWORD=mlflow
+      - POSTGRES_DB=mlflowdb
+    volumes:
+      - ./postgres-temp:/var/lib/postgresql/temp
+
+#  minio:
+#    image: minio/minio:latest
+#    ports:
+#      - "9000:9000"
+#    environment:
+#      - MINIO_ROOT_USER=minioadmin
+#      - MINIO_ROOT_PASSWORD=minioadmin
+#    volumes:
+#      - ./minio-temp:/temp
+#    command: server /temp --console-address ":9001"
+
+  nginx:
+    image: nginx:latest
+    ports:
+      - "80:80"
+    volumes:
+#      - ./nginx.conf:/etc/nginx/nginx.conf:ro
+       - ./nginx.conf:/etc/nginx/conf.d/default.conf:ro
+    depends_on:
+      - mlflow
+#      - minio

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "pad_token_id": 1,
+  "transformers_version": "4.51.3"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2dc6c386af412163c51f18f97152117040b6464f9e64159ef464d50471ceda1c
+size 1101168184

nginx.conf

@@ -0,0 +1,41 @@
+# mlflow_proxy.conf
+
+upstream mlflow {
+    server mlflow:5000;
+}
+
+# upstream minio {
+#     server minio:9000;
+# }
+
+server {
+    listen 80;
+    server_name _; # Opcional: escuta em qualquer nome de host
+
+    # Logs específicos para este server block (ajuda na depuração)
+    access_log /var/log/nginx/mlflow_access.log;
+    error_log /var/log/nginx/mlflow_error.log debug; # Use 'debug' para mais detalhes
+
+    location / {
+        proxy_pass http://mlflow;
+        proxy_set_header Host $host;
+        proxy_set_header X-Real-IP $remote_addr;
+        # Headers úteis para proxys reversos
+        proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+        proxy_set_header X-Forwarded-Proto $scheme;
+
+        # Opcional: Aumentar timeouts se houver problemas de conexão
+        # proxy_connect_timeout 60s;
+        # proxy_read_timeout 60s;
+    }
+
+#     location /minio {
+#         # Atenção: MinIO pode precisar de reescrita de URL ou configuração específica
+#         # dependendo de como ele lida com subpastas.
+#         proxy_pass http://minio;
+#         proxy_set_header Host $host; # MinIO pode precisar do host correto
+#         proxy_set_header X-Real-IP $remote_addr;
+#         proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+#         proxy_set_header X-Forwarded-Proto $scheme;
+#     }
+}

pyproject.toml

@@ -0,0 +1,47 @@
+[project]
+name = "tynerox"
+version = "0.1.0"
+description = ""
+authors = [
+    {name = "Thiago L. Rodrigues"}
+]
+readme = "README.md"
+requires-python = ">=3.9,<4.0"
+dependencies = [
+    "torch (>=2.6.0,<3.0.0)",
+    "transformers[torch] (>=4.50.3,<5.0.0)",
+    "python-dotenv (>=1.1.0,<2.0.0)",
+    "tavily-python (>=0.5.3,<0.6.0)",
+    "langchain-community (>=0.3.20,<0.4.0)",
+    "pydantic (>=2.11.1,<3.0.0)",
+    "pandas (>=2.2.3,<3.0.0)",
+    "openai-agents (>=0.0.7,<0.0.8)",
+    "datasets (>=3.5.0,<4.0.0)",
+    "mlflow (>=2.21.3,<3.0.0)",
+    "beautifulsoup4 (>=4.13.3,<5.0.0)",
+    "packaging (>=24.2,<25.0)",
+    "boto3 (>=1.37.37,<2.0.0)",
+    "flash-attn (>=2.7.4.post1,<3.0.0)",
+]
+
+[tool.poetry]
+name = "tynerox"
+version = "0.1.0"
+packages = [
+  { include = "tynerox", from = "src" }
+]
+
+
+[build-system]
+requires = ["poetry-core>=2.0.0,<3.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.group.dev.dependencies]
+pytest = "^8.3.5"
+
+[tool.pytest.ini_options]
+minversion = "6.0"
+addopts = "-ra -q"
+testpaths = ["src/tests"]
+python_files = ["test_*.py"]
+norecursedirs = ["postgres-data"]

setup.py

@@ -0,0 +1,18 @@
+from setuptools import setup, find_packages
+
+setup(
+    name="tynerox",
+    version="0.1.0",
+    description="TyneRox: custom rotary-transformer causal LM",
+    author="Thiago Luiz Rodrigues",
+    author_email="<EMAIL>",
+    url="https://github.com/seu-usuario/tynerox",
+    license="Apache-2.0",
+    packages=find_packages("src"),
+    package_dir={"": "src"},
+    install_requires=[
+        "torch>=2.6.0,<3.0.0",
+        "transformers[torch]>=4.50.3,<5.0.0",
+        "flash-attn>=2.7.4.post1,<3.0.0",
+    ],
+)

special_tokens_map.json

@@ -0,0 +1,5 @@
+{
+  "eos_token": "<|endoftext|>",
+  "pad_token": "<|endoftext|>",
+  "unk_token": "[UNK]"
+}

src/dataset/fine_tuning.py

@@ -0,0 +1,160 @@
+from typing import List, Dict, Any, Optional
+import torch
+from torch.utils.data import DataLoader
+from datasets import load_dataset
+from transformers import AutoTokenizer
+from torch.nn.utils.rnn import pad_sequence
+from functools import partial
+
+
+def tokenize_function(examples: Dict[str, Any], tokenizer: Any) -> Dict[str, List[int]]:
+    """
+    Aplica a template de chat do tokenizer e gera os token ids.
+
+    Args:
+        examples (Dict[str, Any]): Dicionário contendo a lista de mensagens sob a chave "messages".
+        tokenizer (Any): Instância do tokenizer que deverá possuir a propriedade 'chat_template'.
+
+    Returns:
+        Dict[str, List[int]]: Dicionário com os token ids gerados.
+    """
+    full_text = tokenizer.apply_chat_template(
+        examples["messages"],
+        tokenize=True,
+        add_generation_prompt=True
+    )
+    return {"input_ids": full_text}
+
+
+def custom_collate_fn(
+    batch: List[Dict[str, List[int]]],
+    pad_token_id: int = 29797,
+    ignore_index: int = -100,
+    allowed_max_length: Optional[int] = None,
+    device: str = "cpu",
+) -> Dict[str, torch.Tensor]:
+    """
+    • Faz padding das sequências
+    • Cria pares (input, label) deslocando 1 posição
+    • Aplica `ignore_index` (-100) APENAS nos labels depois do 1.º PAD
+    """
+
+    # 1) Lista → Tensor  +  PAD final
+    seqs = [torch.tensor(s["input_ids"] + [pad_token_id]) for s in batch]
+
+    # 2) Padding até o comprimento máximo do batch
+    padded = pad_sequence(seqs, batch_first=True, padding_value=pad_token_id)
+
+    # 3) Desloca 1 posição e CLONA para quebrar o compartilhamento de memória
+    input_ids = padded[:, :-1].clone()     # ← nunca terá -100
+    labels    = padded[:, 1:].clone()      # ← vamos editar aqui
+
+    # 4) Define -100 após o primeiro PAD de cada sequência
+    pad_mask = (labels == pad_token_id)
+    if pad_mask.any():
+        # índice da primeira ocorrência de PAD em cada linha
+        first_pad_pos = pad_mask.float().cumsum(1).eq(1) & pad_mask
+        # tudo que vem depois do primeiro PAD recebe -100
+        mask_after_first_pad = pad_mask & ~first_pad_pos
+        labels[mask_after_first_pad] = ignore_index
+
+    # 5) Trunca se for solicitado
+    if allowed_max_length is not None:
+        input_ids = input_ids[:, :allowed_max_length]
+        labels    = labels[:, :allowed_max_length]
+
+    return {
+        "input_ids": input_ids.to(device),
+        "labels":    labels.to(device),
+    }
+
+
+
+def create_data_loader_fine_tuning(
+        tokenizer: Any,
+        batch_size: int,
+        path_folder: str,
+        split: str = "train",
+        pad_token_id: int = 0,
+        ignore_index: int = -100,
+        allowed_max_length: Optional[int] = None,
+        device: str = "cpu"
+) -> DataLoader:
+    """
+    Cria o DataLoader para fine-tuning, a partir de um dataset_files tokenizado.
+
+    Esta função carrega o dataset_files, aplica a tokenização utilizando uma template de chat,
+    e retorna um DataLoader que utiliza a função custom_collate_fn para o processamento
+    adequado das batches.
+
+    Args:
+        tokenizer (Any): Tokenizer pré-treinado que suporte chat templates.
+        batch_size (int): Número de amostras por batch.
+        path_folder (str): Caminho ou identificador do dataset_files.
+        split (str): Divisão do dataset_files a ser utilizada (por exemplo, "train" ou "test").
+        pad_token_id (int): ID do token para padding.
+        ignore_index (int): Valor a ser ignorado na função de perda.
+        allowed_max_length (Optional[int]): Se definido, trunca as sequências para este tamanho máximo.
+        device (str): Dispositivo para onde os tensores serão enviados ("cpu" ou "cuda").
+
+    Returns:
+        DataLoader: Instância do DataLoader pronta para o fine-tuning.
+    """
+    # Define a template de chat e atribui ao tokenizer.
+    chat_template = """
+    {% for message in messages %}
+        {% if message['role'] == 'user' %}
+            {{ '<|user_start|>' + message['content'] + '<|user_end|>' + '\n'}}
+        {% elif message['role'] == 'assistant' %}
+            {{ '<|assistant_start|>' + message['content'] + '<|assistant_end|>' + '\n' }}
+        {% endif %}
+    {% endfor %}
+    """
+    tokenizer.chat_template = chat_template
+
+    # Carrega o dataset_files.
+    raw_dataset = load_dataset(path=path_folder, split=split, download_mode="force_redownload")
+
+    # Aplica a tokenização utilizando a função definida.
+    tokenized_dataset = raw_dataset.map(
+        lambda examples: tokenize_function(examples, tokenizer),
+        batched=True,
+        remove_columns=raw_dataset.column_names,
+        desc="Tokenizando dataset_files"
+    )
+
+    # Configura o collate_fn com os parâmetros desejados.
+    collate = partial(
+        custom_collate_fn,
+        pad_token_id=pad_token_id,
+        ignore_index=ignore_index,
+        allowed_max_length=allowed_max_length,
+        device=device
+    )
+
+    print("Criando DataLoader...")
+    return DataLoader(
+        tokenized_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        drop_last=False,
+        num_workers=0,
+        collate_fn=collate
+    )
+
+
+if __name__ == "__main__":
+    # Carrega o tokenizer pré-treinado.
+    tokenizer = AutoTokenizer.from_pretrained("neuralmind/bert-base-portuguese-cased")
+
+    # Cria o DataLoader para a divisão de treino do dataset_files "conversational".
+    loader = create_data_loader_fine_tuning(
+        tokenizer=tokenizer,
+        batch_size=100,
+        path_folder="conversational",
+        split="test"
+    )
+
+    # Testa a extração de uma batch.
+    batch = next(iter(loader))
+    print(batch["input_ids"].shape, batch["labels"].shape)

src/dataset/pre_train.py

@@ -0,0 +1,106 @@
+import os
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast
+from datasets import load_dataset, Dataset, DatasetDict
+
+
+def tokenize_function(
+    examples: Dict[str, List[str]],
+    tokenizer: PreTrainedTokenizer
+) -> Dict[str, List[List[int]]]:
+    """
+    Tokeniza os exemplos sem aplicar truncamento ou padding.
+
+    Retorna apenas os input_ids.
+    """
+    tokenized_output = tokenizer(examples["text"], truncation=False, padding=False)
+    return {"input_ids": tokenized_output["input_ids"]}
+
+
+def pack_documents(
+    examples: Dict[str, List[List[int]]],
+    max_length: int,
+    eos_token_id: Optional[int] = None
+) -> Dict[str, List[List[int]]]:
+    """
+    Aplica Document Packing e retorna apenas os inputs de tamanho fixo (max_length),
+    descartando o último token extra usado para labels.
+    """
+    # Concatena tokens de todo o batch
+    concatenated: List[int] = []
+    separator = [eos_token_id] if eos_token_id is not None else []
+    first = True
+    for doc in examples["input_ids"]:
+        if not first and separator:
+            concatenated.extend(separator)
+        concatenated.extend(doc)
+        first = False
+
+    block_size = max_length + 1
+    total_len = (len(concatenated) // block_size) * block_size
+    if total_len == 0:
+        return {"input_ids": []}
+
+    concatenated = concatenated[:total_len]
+    # Divide em blocos de block_size e remove o último token de cada bloco
+    blocks = [
+        concatenated[i : i + block_size]
+        for i in range(0, total_len, block_size)
+    ]
+    inputs = [blk[:-1] for blk in blocks]
+
+    # Filtra qualquer bloco vazio
+    inputs = [inp for inp in inputs if len(inp) > 0]
+    return {"input_ids": inputs}
+
+
+def create_train_dataloader(
+    folder_path: str,
+    tokenizer: PreTrainedTokenizerFast,
+    batch_size: int = 4,
+    max_length: int = 512,
+    drop_last: bool = True,
+    num_workers: int = 5
+) -> Optional[DataLoader]:
+    """
+    Carrega .txt de folder_path, tokeniza, aplica packing só de inputs
+    e retorna um DataLoader que fornece batches de input_ids.
+    """
+    raw_dataset = load_dataset(folder_path, split="train", streaming=False)
+    print(f"Dataset bruto carregado: {raw_dataset}")
+
+    # 1) Tokenização
+    tokenized = raw_dataset.map(
+        lambda ex: tokenize_function(ex, tokenizer),
+        batched=True,
+        batch_size=1000,
+        num_proc=20,
+        remove_columns=raw_dataset.column_names,
+    )
+    print(f"Dataset tokenizado: {tokenized}")
+
+    # 2) Document Packing sem labels
+    packed = tokenized.map(
+        lambda ex: pack_documents(
+            ex,
+            max_length=max_length,
+            eos_token_id=tokenizer.eos_token_id
+        ),
+        batched=True,
+        batch_size=10000,
+        num_proc=20,
+    )
+
+    # 3) Configura para PyTorch
+    packed.set_format(type="torch", columns=["input_ids"])
+
+    print("Criando DataLoader...")
+    return DataLoader(
+        packed,
+        batch_size=batch_size,
+        drop_last=drop_last,
+        num_workers=num_workers,
+    )

src/logger/logger.py

@@ -0,0 +1,144 @@
+from typing import Dict, Optional
+import os
+from zoneinfo import ZoneInfo
+
+import mlflow
+import pandas as pd
+import torch
+import torch.nn as nn
+from datetime import datetime, date
+
+
+class TrainerLogger:
+    def __init__(
+        self,
+        tracking_uri: str,
+        experiment: str,
+        total_params: int,
+        model_name: str = None,
+        run_name: str = None,
+        tags: Dict[str, str] = None,
+    ):
+        mlflow.set_tracking_uri(tracking_uri)
+        mlflow.set_experiment(experiment)
+
+        # Ativar autologging para PyTorch
+        mlflow.pytorch.autolog(log_models=True)  # Desativamos log automático de modelos para controle manual
+
+        # Iniciar run com contexto
+        self.run = mlflow.start_run(run_name=run_name)
+        self.run_id = self.run.info.run_id
+        self.experiment = experiment
+        self.model_name = model_name
+        self.total_params = total_params
+
+        # Registrar tags para melhor organização
+        default_tags = {"model_type": self.model_name}
+        if tags:
+            default_tags.update(tags)
+        mlflow.set_tags(default_tags)
+
+        # Registrar parâmetros
+        base_params = {"model_name": self.model_name, "total_params": self.total_params}
+        self.log_parameters(base_params)
+
+    def log_parameters(self, parameters: dict):
+        mlflow.log_params(parameters)  # Mais eficiente que log_param individual
+
+    def log_metrics(self, metrics: dict, step: Optional[int] = None):
+        mlflow.log_metrics(metrics, step)
+
+    def log_checkpoint_table(self, current_lr:float, loss:float, perplexity: float, last_batch:int) -> None:
+        """
+        Log a checkpoint record (month, day, hour, perplexity) to MLflow as a table artifact.
+        Perplexity is rounded to 4 decimal places.
+
+        Parameters
+        ----------
+        perplexity : float
+            The perplexity metric to log (rounded to 4 decimal places).
+            :param current_lr:
+            :param loss:
+            :param perplexity:
+            :param last_batch:
+        """
+        # Define artifact directory and ensure it exists
+        artifact_dir = f"checkpoint_table/model"
+        os.makedirs(artifact_dir, exist_ok=True)
+
+        # Capture current timestamp
+        now = datetime.now(ZoneInfo("America/Sao_Paulo"))
+        record = {
+            "month": now.month,
+            "day": now.day,
+            "hour": f"{now.hour:02d}:{now.minute:02d}",
+            "last_batch": last_batch,
+            "current_lr": round(current_lr, 7),
+            "perplexity": round(perplexity, 4),
+            "loss": round(loss, 4),
+
+        }
+        df_record = pd.DataFrame([record])
+
+        # Define artifact file path (relative POSIX path)
+        artifact_file = f"{artifact_dir}/checkpoint_table.json"
+
+        # Log the table to MLflow Tracking
+        mlflow.log_table(
+            data=df_record,
+            artifact_file=artifact_file
+        )
+
+    def checkpoint_model(self, model: nn.Module):
+        # Criar diretório local para checkpoint
+        step = 1
+        checkpoint_dir = f"checkpoints/model_{step}"
+        os.makedirs(checkpoint_dir, exist_ok=True)
+
+        # Salvar estado do modelo localmente
+        checkpoint_path = os.path.join(checkpoint_dir, "model.pth")
+        torch.save(model.state_dict(), checkpoint_path)
+
+        # Registrar artefato no MLflow
+        mlflow.log_artifact(checkpoint_path, f"model_checkpoints/epoch_{step}")
+
+        input_example = torch.zeros(1, 128, dtype=torch.long)  # Ajuste as dimensões conforme seu modelo
+        # input_example_numpy = input_example.cpu().numpy()
+
+        # Registrar modelo no registro de modelos MLflow
+        if self.model_name:
+            registered_model_name = f"{self.model_name}"
+            mlflow.pytorch.log_model(
+                pytorch_model=model,
+                artifact_path=f"models/epoch_{step}",
+                registered_model_name=registered_model_name,
+                pip_requirements=["torch>=1.9.0"],
+                code_paths=["tynerox/"],  # Inclui código-fonte relevante
+                # input_example=input_example_numpy,  # Exemplo de entrada
+                signature=None  # Adicione assinatura do modelo se possível
+            )
+
+        table_dict = {
+            "entrada": ["Pergunta A", "Pergunta B"],
+            "saida": ["Resposta A", "Resposta B"],
+            "nota": [0.75, 0.40],
+        }
+
+    def log_html(self, html: str, step: Optional[int] = None):
+        file_path = f"visualizations/sample.html"
+        os.makedirs(os.path.dirname(file_path), exist_ok=True)
+
+        with open(file_path, "w") as f:
+            f.write(html)
+
+        mlflow.log_artifact(file_path)
+
+    def finish(self):
+        """Finaliza a execução do MLflow run"""
+        mlflow.end_run()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.finish()

src/pre-training.py

@@ -0,0 +1,111 @@
+import math
+
+import torch
+from tokenizers import Tokenizer
+from transformers import PreTrainedTokenizerFast, get_cosine_schedule_with_warmup
+
+from training import PreTrainer
+from tynerox.modeling import TyneRoxModel, TyneRoxConfig
+from dataset.pre_train import create_train_dataloader
+
+if __name__ == "__main__":
+
+    # 1 - Carrega o tokenizador
+    tokenizer = Tokenizer.from_file("tokenizer/tokens-bpe-36k.json")
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer,
+        unk_token="[UNK]",
+        pad_token="<|endoftext|>",
+        eos_token="<|endoftext|>",
+    )
+
+    tokenizer.save_pretrained(f"../")
+
+    # 2 Inicia a configuração e o modelo
+    config = TyneRoxConfig(
+        vocab_size=tokenizer.vocab_size,
+        pad_token_id=tokenizer.pad_token_id,
+    )
+
+    model = TyneRoxModel(config)
+    model.to("cuda")
+
+    # 3 - Carrega o dataset de treinamento
+    folder_path = "bobboyms/subset-Itau-Unibanco-aroeira-1B-tokens"
+    dataloader = create_train_dataloader(
+        folder_path,
+        tokenizer,
+        batch_size=20,
+        max_length=1024,
+        drop_last=True,
+        num_workers=10
+    )
+
+    # 4 - Criando o optmizer
+    model = torch.compile(model)
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=0.000461,  # Mantenha a LR inicial ou ajuste ligeiramente (ex: 3e-4)
+        weight_decay=0.1
+    )
+
+    # 5 - Configura o warmup
+    epochs = 1
+    batch_size = 40
+    size_dataset = 2_883_231
+    warmup_ratio = 0.05
+
+    num_training_steps = len(dataloader) * epochs
+    num_warmup_steps = math.floor(num_training_steps * warmup_ratio)
+
+    # 6. Scheduler
+    scheduler = get_cosine_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=num_warmup_steps,
+        num_training_steps=num_training_steps,
+    )
+
+    sample_prompts = [
+        "Olá, como vai você? ",
+        "Quando a manhã chegou, Iracema ainda estava ali, debruçada, como uma borboleta que ",
+        "Não, respondeu; na verdade, estou com medo ",
+        "O resultado representa uma desaceleração ",
+        "No vídeo, é possível ver ",
+        "Essa receita de torta de frango ",
+        "Durante o primeiro mandato ",
+        "Os donos de cães "
+    ]
+
+    logger_config = {
+        "tracking_uri": "http://127.0.0.1:5000",
+        "experiment": "Pre training LLM",
+        "model_name": "Pre training LLM (Long Context)"
+    }
+
+    trainer = PreTrainer(
+        model=model,
+        optimizer=optimizer,
+        scheduler=scheduler,
+        tokenizer=tokenizer,
+        train_loader=dataloader,
+        test_loader=None,
+        logger_config=logger_config,
+        use_amp=True
+    )
+
+    trainer.train(num_epochs=epochs,sample_prompts=sample_prompts)
+
+    # 7 - Salva as configurações do modelo para enviar para o hugginfaces
+    model.save_pretrained(f"../")
+
+
+
+
+
+
+
+
+
+
+
+

src/tokenizer/tests.py

@@ -0,0 +1,77 @@
+from datasets import load_dataset
+from tokenizers import Tokenizer
+
+if __name__ == "__main__":
+    # Carrega streaming do dataset e o tokenizer
+    dataset_stream = load_dataset(
+        "bobboyms/subset-Itau-Unibanco-aroeira-1B-tokens",
+        split="train",
+        streaming=True
+    )
+    tokenizer = Tokenizer.from_file("tokens-bpe-36k.json")
+    encode    = tokenizer.encode
+    unk_id    = tokenizer.token_to_id("[UNK]")
+    vocab_size = tokenizer.get_vocab_size()
+
+    print("Tamanho do vocabulário:", tokenizer.get_vocab_size())
+    enc = tokenizer.encode("Apostas combinadas: Fantástico exibe mensagens exclusivas da investigação contra Bruno Henrique, do Flamengo")
+    print(tokenizer.decode(enc.ids, skip_special_tokens=True))
+
+    # Contadores
+    total_tokens = 0
+    total_words  = 0
+    unk_tokens   = 0
+    seen_ids     = set()
+
+    batch_size   = 512
+    batch_counter = 0
+
+    def batch_iterator(stream, bs):
+        buf = []
+        for ex in stream:
+            buf.append(ex["text"])
+            if len(buf) == bs:
+                yield buf
+                buf = []
+        if buf:
+            yield buf
+
+    for texts in batch_iterator(dataset_stream, batch_size):
+        # tokeniza em batch
+        encs = tokenizer.encode_batch(texts)
+
+        # conta palavras e tokens no batch
+        words_in_batch = sum(len(t.split()) for t in texts)
+        total_words   += words_in_batch
+
+        for enc in encs:
+            total_tokens += len(enc.ids)
+            unk_tokens   += enc.ids.count(unk_id)
+            seen_ids.update(enc.ids)
+
+        # impressão parcial a cada 100 batches
+        if batch_counter % 100 == 0:
+            oov_rate  = unk_tokens / total_tokens * 100
+            frag      = total_tokens / total_words
+            coverage  = len(seen_ids) / vocab_size * 100
+            ttr       = len(seen_ids) / total_tokens
+            print(f"[Batch {batch_counter:04d}] "
+                  f"OOV: {oov_rate:.3f}% | "
+                  f"Frag: {frag:.3f} t/palavra | "
+                  f"Coverage: {coverage:.2f}% | "
+                  f"TTR: {ttr:.4f}")
+        batch_counter += 1
+
+    # resultado final
+    oov_rate  = unk_tokens / total_tokens * 100
+    frag      = total_tokens / total_words
+    coverage  = len(seen_ids) / vocab_size * 100
+    ttr       = len(seen_ids) / total_tokens
+
+    print("\n=== Métricas Finais ===")
+    print(f"Total de tokens: {total_tokens}")
+    print(f"Total de palavras: {total_words}")
+    print(f"OOV rate:         {oov_rate:.3f}%")
+    print(f"Fragmentação:     {frag:.3f} tokens/palavra")
+    print(f"Voc. coverage:    {coverage:.2f}% do vocabulário usado")
+    print(f"Type–Token Ratio: {ttr:.4f}")

src/tokenizer/trainer.py

@@ -0,0 +1,77 @@
+from datasets import load_dataset
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace, ByteLevel
+import time  # Para medir o tempo
+from tokenizers.normalizers import Sequence, NFD, Lowercase, StripAccents, NFC
+from tokenizers.decoders import ByteLevel as ByteLevelDecoder
+# 1. Carregar o dataset em modo streaming
+dataset_stream = load_dataset("bobboyms/subset-Itau-Unibanco-aroeira-1B-tokens", split="train", streaming=True)
+
+print("Dataset carregado em modo streaming:")
+print(dataset_stream)
+
+# Nome da coluna que contém o texto
+coluna_texto = "text"
+
+
+# 2. Criar o gerador para o treinamento do tokenizador
+# Esta função irá iterar sobre o dataset streaming e fornecer o texto
+def get_training_corpus_streaming():
+    count = 0
+    start_time = time.time()
+    print("Iniciando iteração sobre o dataset streaming para o tokenizador...")
+    for sample in dataset_stream:
+        # Certifique-se de que a amostra não é None e a coluna existe
+        if sample and coluna_texto in sample and isinstance(sample[coluna_texto], str):
+            yield sample[coluna_texto]
+            count += 1
+            if count % 10000 == 0:  # Log a cada 10000 amostras
+                elapsed = time.time() - start_time
+                print(f"  Processadas {count} amostras para o tokenizador em {elapsed:.2f} segundos...")
+        else:  # Opcional: Logar amostras inválidas/puladas
+            print(f"Aviso: Pulando amostra inválida ou sem coluna '{coluna_texto}': {sample}")
+    end_time = time.time()
+    print(
+        f"Iteração completa. Total de {count} amostras fornecidas ao tokenizador em {end_time - start_time:.2f} segundos.")
+
+special_tokens=[
+            "[UNK]", "<|endoftext|>",
+            "<|user_start|>", "<|user_end|>",
+            "<|assistant_start|>", "<|assistant_end|>",
+            "<|think_start|>", "<|think_end|>",
+            "<|command_start|>", "<|command_end|>",
+        ]
+
+if __name__ == "__main__":
+    print("Inicializando o tokenizador BPE...")
+    # tokenizer.pre_tokenizer = Whitespace()
+    tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+    tokenizer.pre_tokenizer = ByteLevel(add_prefix_space=True)
+    tokenizer.normalizer = NFC()
+    tokenizer.decoder = ByteLevelDecoder(add_prefix_space=True)
+
+    # Aqui: merges mais agressivos
+    trainer = BpeTrainer(
+        vocab_size=36000 + len(special_tokens),
+        min_frequency=7,
+        limit_alphabet=1300,
+        # continuing_subword_prefix="##",
+        # end_of_word_suffix="</w>",               # baixa frequência mínima para 1
+        show_progress=True,             # barra de progresso
+        special_tokens=special_tokens,
+    )
+
+    print("Iniciando o treinamento do tokenizador a partir do stream...")
+    start_train_time = time.time()
+    tokenizer.train_from_iterator(
+        get_training_corpus_streaming(),
+        trainer=trainer
+    )
+    end_train_time = time.time()
+    print(f"Treinamento do tokenizador concluído em {end_train_time - start_train_time:.2f} segundos!")
+
+    save_path = "tokens-bpe-36k.json"
+    tokenizer.save("tokens-bpe-36k.json", pretty=True)
+    print(f"Tokenizador salvo em {save_path}")

src/training.py

@@ -0,0 +1,203 @@
+import math
+import time
+from typing import Any, Optional, Dict, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from logger.logger import TrainerLogger
+from torch.utils.data import DataLoader
+from transformers import PreTrainedModel
+
+# Configuração do dispositivo
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class BaseTrainer:
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        optimizer: torch.optim.Optimizer,
+        scheduler: torch.optim.lr_scheduler._LRScheduler,
+        tokenizer: Any,
+        train_loader: DataLoader,
+        test_loader: Optional[DataLoader] = None,
+        logger_config: Dict[str, Any] = None,
+        use_amp: bool = True,
+    ):
+        self.model = model.to(device)
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.tokenizer = tokenizer
+        self.train_loader = train_loader
+        self.test_loader = test_loader
+        self.use_amp = use_amp
+        self.scaler = torch.amp.GradScaler('cuda') if use_amp else None
+        self.train_step = 0
+        self._best_perplexity = float('inf')
+        self._epochs_no_improve = 0
+
+        total_params = sum(p.numel() for p in model.parameters())
+        self.logger = TrainerLogger(
+            tracking_uri=logger_config["tracking_uri"],
+            experiment=logger_config["experiment"],
+            run_name=logger_config["model_name"],
+            model_name=logger_config["model_name"],
+            total_params=total_params,
+            tags={"version": "1.0", "environment": "development"},
+        )
+
+    def _generate_sample(self, sample_prompts: List[str] = []):
+        self.model.eval()
+        samples_html = ""
+        for prompt in sample_prompts:
+            try:
+                # sample_text = generate_text_sample(self.model, self.tokenizer, prompt)
+                inputs = self.tokenizer(prompt, return_tensors="pt")
+                input_ids = inputs.input_ids.to(self.model.device)
+
+                # 4) Gere texto
+                with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
+                    generated_ids = self.model.generate(
+                        input_ids=input_ids,
+                        max_length=100,  # comprimento total (prompt + continuação)
+                        num_beams=5,  # número de “hips” em beam search
+                        do_sample=True,  # ativa amostragem (em vez de pura greed)
+                        top_k=50,  # restringe sampling aos top-50 tokens
+                        top_p=0.95,  # usa nucleus sampling (p acumulado ≤ 0.95)
+                        temperature=0.7,  # controle de “criatividade”
+                        repetition_penalty=1.2,  # penaliza repetições exatas
+                        use_cache=True,  # reutiliza past_key_values (default)
+                        eos_token_id=self.tokenizer.eos_token_id,
+                        pad_token_id=self.tokenizer.pad_token_id,
+                    )
+
+                # 5) Decode para string
+                generated_text = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+            except Exception as e:
+                generated_text = f"Erro: {e}"
+            samples_html += f"<h4><b>prompt:</b> {prompt}</h4><p><b>Resposta:</b> {generated_text}</p>"
+        self.model.train()
+        return samples_html
+
+    def _calc_loss_batch(self, inputs: torch.Tensor) -> torch.Tensor:
+        """
+        Calcula apenas a entropia cruzada para um batch de input_ids,
+        desativando o cache de chaves/valores durante o treinamento.
+        """
+        ignore_idx = -100
+        # valida que todos os tokens estão no vocabulário ou são tokens de ignore
+        valid = ((inputs >= 0) | (inputs == ignore_idx)) & (inputs < self.tokenizer.vocab_size)
+        assert valid.all(), f"Há labels inválidos: min={inputs.min().item()}, max={inputs.max().item()}"
+
+        inputs = inputs.to(device)
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            outputs = self.model(
+                input_ids=inputs,
+                labels=inputs,
+                use_cache=False,  # desabilita o KV-cache no treino
+                return_dict=True  # garante acesso via .loss e .logits
+            )
+            loss = outputs.loss
+            logits = outputs.logits
+            if torch.isnan(logits).any() or torch.isinf(logits).any():
+                raise RuntimeError("Logits inválidos detectados")
+        return loss
+
+    def _train_epoch(self, epoch: int, sample_prompts: Optional[List[str]] = None) -> List[float]:
+        if sample_prompts is None:
+            sample_prompts = []
+
+        self.model.train()
+        losses = []
+        size_dataset = len(self.train_loader)
+        pbar = tqdm(
+            self.train_loader,
+            total=size_dataset,
+            desc=f"Epoch {epoch + 1}",
+            unit="batch",
+            leave=False,
+        )
+
+        for i, batch in enumerate(pbar):
+            start_time = time.time()
+            self.optimizer.zero_grad()
+            loss = self._calc_loss_batch(batch['input_ids'])
+            losses.append(loss.item())
+
+            if self.use_amp:
+                self.scaler.scale(loss).backward()
+                self.scaler.unscale_(self.optimizer)
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+                self.scaler.step(self.optimizer)
+                self.scaler.update()
+            else:
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+                self.optimizer.step()
+
+            self.scheduler.step()
+            perplexity = math.exp(loss.item())
+            current_lr = self.optimizer.param_groups[0].get('lr', 0.0)
+            elapsed_time = time.time() - start_time
+
+            pbar.set_postfix({
+                "loss": f"{loss.item():.4f}",
+                "perplexity": f"{perplexity:.4f}",
+                "lr": f"{current_lr:.2e}",
+                "elapsed_time": f"{elapsed_time:.2f}s",
+            })
+
+            # Logging a cada 100 batches
+            if (i + 1) % 100 == 0:
+                self.train_step += 1
+                avg_loss = sum(losses[-100:]) / 100
+                avg_perplexity = math.exp(sum(losses[-100:]) / 100)
+                self.logger.log_metrics(
+                    {
+                        "train_loss": avg_loss,
+                        "train_perplexity": avg_perplexity,
+                        "lr": current_lr,
+                    },
+                    step=self.train_step,
+                )
+
+            # Gera samples
+            if (i + 1) % 500 == 0:
+                samples_html = self._generate_sample(sample_prompts)
+                self.logger.log_html(f"<html><head><meta charset='utf-8'></head><body>{samples_html}</body></html>",
+                                     step=self.train_step)
+
+            # Checkpoint a cada 1000 batches
+            if (i + 1) % 1000 == 0:
+                avg_loss = sum(losses[-1000:]) / 1000
+                avg_perplexity = math.exp(sum(losses[-1000:]) / 1000)
+                self.logger.log_checkpoint_table(current_lr, avg_loss, avg_perplexity, i + 1)
+                self.logger.checkpoint_model(self.model)
+                self.model.save_pretrained(f"../")
+
+
+        return losses
+
+    def train(self, num_epochs: int = 500, sample_prompts: Optional[List[str]] = None):
+        for epoch in range(num_epochs):
+            train_losses = self._train_epoch(epoch, sample_prompts)
+            mean_train_loss = sum(train_losses) / len(train_losses)
+            self.logger.log_metrics(
+                {"mean_train_loss": mean_train_loss},
+                step=epoch,
+            )
+            print(f"Epoch {epoch + 1} | Train Loss: {mean_train_loss:.4f}")
+
+        self.logger.finish()
+        print("Treinamento concluído!")
+
+
+# Exemplo de uso para Fine-Tuning:
+class TuningTrainer(BaseTrainer):
+    pass
+
+# Exemplo de uso para Pré-Treinamento:
+class PreTrainer(BaseTrainer):
+    pass

src/tynerox/__init__.py

@@ -0,0 +1,3 @@
+from .modeling import TyneRoxModel, TyneRoxConfig
+
+__all__ = ["TyneRoxConfig", "TyneRoxModel"]

src/tynerox/modeling.py

@@ -0,0 +1,449 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from flash_attn.flash_attn_interface import flash_attn_func
+from transformers import PretrainedConfig, PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from typing import Optional, Literal, Union, Tuple
+
+
+class PositionalEncoding(nn.Module):
+    """
+    Implements positional encoding (sinusoidal or rotary).
+    """
+    def __init__(
+        self,
+        embed_dim: int,
+        context_length: int,
+        dropout: float = 0.1,
+        encoding_type: Literal['sinusoidal', 'rotary'] = 'rotary',
+    ):
+        super().__init__()
+        if embed_dim <= 0 or context_length <= 0:
+            raise ValueError("embed_dim and context_length must be positive integers")
+        if not 0 <= dropout < 1:
+            raise ValueError("dropout must be between 0 and 1")
+
+        self.dropout = nn.Dropout(dropout)
+        self.encoding_type = encoding_type.lower()
+        self.max_seq_len = context_length
+        self.embed_dim = embed_dim
+
+        if self.encoding_type == 'sinusoidal':
+            pe = self._create_sinusoidal_embeddings(context_length, embed_dim)
+            self.register_buffer('pe', pe.unsqueeze(0), persistent=True)
+        elif self.encoding_type == 'rotary':
+            if embed_dim % 2 != 0:
+                raise ValueError("embed_dim must be even for rotary encoding")
+            # inv_freq of size D/2
+            inv_freq = 1.0 / (10000 ** (torch.arange(0, embed_dim, 2).float() / embed_dim))
+            self.register_buffer('inv_freq', inv_freq, persistent=True)
+        else:
+            raise ValueError("Unsupported encoding_type: 'sinusoidal' or 'rotary'")
+
+    def _create_sinusoidal_embeddings(self, max_seq_len: int, dim: int) -> torch.Tensor:
+        position = torch.arange(max_seq_len).unsqueeze(1).float()
+        div_term = torch.exp(torch.arange(0, dim, 2).float() * (-math.log(10000.0) / dim))
+        pe = torch.zeros(max_seq_len, dim)
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        return pe
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # x shape: [B, T, D]
+        if self.encoding_type == 'sinusoidal':
+            seq_len = x.size(1)
+            x = x + self.pe[:, :seq_len, :]
+        else:
+            # rotary: split even/odd dims and apply rotary
+            seq_len = x.size(1)
+            positions = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+            # freqs of shape [T, D/2]
+            freqs = torch.einsum('i , j -> i j', positions, self.inv_freq)
+            x = self.apply_rotary(x, freqs)
+        return self.dropout(x)
+
+    @staticmethod
+    def apply_rotary(x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
+        # x: [B, T, D], emb: [T, D/2]
+        x1, x2 = x.chunk(2, dim=-1)              # each [B, T, D/2]
+        emb_sin = emb.sin()[None, :, :]          # [1, T, D/2]
+        emb_cos = emb.cos()[None, :, :]          # [1, T, D/2]
+        # apply rotary
+        rotated1 = x1 * emb_cos + x2 * emb_sin
+        rotated2 = x2 * emb_cos - x1 * emb_sin
+        return torch.cat([rotated1, rotated2], dim=-1)  # [B, T, D]
+
+
+class PositionalEmbedding(nn.Module):
+    """
+    Combines token embedding with positional encoding.
+    """
+    def __init__(
+        self,
+        vocab_size: int,
+        embed_dim: int,
+        context_length: int,
+        dropout: float = 0.05,
+        encoding_type: Literal['sinusoidal', 'rotary'] = 'rotary'
+    ):
+        super().__init__()
+        if vocab_size <= 0 or embed_dim <= 0 or context_length <= 0:
+            raise ValueError("vocab_size, embed_dim, context_length must be > 0")
+
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.scale = math.sqrt(embed_dim)
+        self.pos_encoding = PositionalEncoding(
+            embed_dim=embed_dim,
+            context_length=context_length,
+            dropout=dropout,
+            encoding_type=encoding_type
+        )
+
+    def forward(self, input_ids: torch.LongTensor) -> torch.Tensor:
+        # input_ids: [B, T]
+        x = self.token_embedding(input_ids) * self.scale  # [B, T, D]
+        return self.pos_encoding(x)
+
+
+def get_alibi_slopes(n_heads: int) -> torch.Tensor:
+    def _get_slopes(n):
+        base = 2 ** (-8.0 / n)
+        return torch.tensor([base ** (i + 1) for i in range(n)])
+    if math.log2(n_heads).is_integer():
+        return _get_slopes(n_heads)
+    m = 2 ** math.floor(math.log2(n_heads))
+    slopes = _get_slopes(m)
+    extra = _get_slopes(2 * m)[::2][: n_heads - m]
+    return torch.cat([slopes, extra], dim=0)
+
+# -----------------------------------------------------------------------------
+# Feed-Forward
+# -----------------------------------------------------------------------------
+
+class FeedForward(nn.Module):
+    def __init__(self, emb_dim: int, hidden_dim_multiplier: int = 4):
+        super().__init__()
+        hidden_dim = emb_dim * hidden_dim_multiplier
+        self.fc1 = nn.Linear(emb_dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim // 2, emb_dim)
+        self.activation = nn.SiLU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_fc1 = self.fc1(x)
+        x_up, x_gate = x_fc1.chunk(2, dim=-1)
+        return self.fc2(x_up * self.activation(x_gate))
+
+# -----------------------------------------------------------------------------
+# Attention-Free Transformer (AFT) Simple
+# -----------------------------------------------------------------------------
+
+class AFTSimple(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        activation=torch.sigmoid,
+        causal: bool = True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.causal = causal
+        self.activation = activation
+
+        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=False)
+        self.project = nn.Linear(embed_dim, embed_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # x: [B, T_new, D]
+        B, T_new, D = x.shape
+        if D != self.embed_dim:
+            raise ValueError(f"Input dim ({D}) != embed_dim ({self.embed_dim})")
+
+        qkv = self.qkv(x)                 # [B, T_new, 3*D]
+        Q, K_new, V_new = qkv.chunk(3, dim=-1)  # each [B, T_new, D]
+
+        # concatenate past if provided
+        if past_key_values is not None:
+            K_past, V_past = past_key_values
+            K = torch.cat([K_past, K_new], dim=1)  # [B, T_all, D]
+            V = torch.cat([V_past, V_new], dim=1)
+        else:
+            K, V = K_new, V_new
+
+        # compute attention-free aggregate
+        softmax_k = F.softmax(K, dim=1)          # [B, T_all, D]
+        weighted_v = softmax_k * V               # [B, T_all, D]
+
+        if self.causal:
+            context = torch.cumsum(weighted_v, dim=1)  # [B, T_all, D]
+        else:
+            total = weighted_v.sum(dim=1, keepdim=True)  # [B, 1, D]
+            context = total.expand(-1, K.size(1), -1)    # [B, T_all, D]
+
+        # slice only the new positions
+        context_new = context[:, -T_new:, :]     # [B, T_new, D]
+        gate = self.activation(Q)                # [B, T_new, D]
+        Y = gate * context_new                   # [B, T_new, D]
+        Y = self.project(Y)                      # [B, T_new, D]
+
+        # return output and updated cache
+        return Y, (K, V)
+
+# -----------------------------------------------------------------------------
+# Flash Attention with ALiBi and KV-cache
+# -----------------------------------------------------------------------------
+
+class FlashAttention(nn.Module):
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        window_size: int,
+        causal: bool = True,
+        qkv_bias: bool = False,
+    ):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+        self.causal = causal
+        self.window_size = window_size
+
+        self.qkv = nn.Linear(embed_dim, 3 * embed_dim, bias=qkv_bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=qkv_bias)
+
+        # precompute ALiBi slopes
+        self.register_buffer('alibi', get_alibi_slopes(num_heads))
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        # x: [B, T_new, D]
+        B, T_new, _ = x.size()
+        qkv = self.qkv(x).view(B, T_new, self.num_heads, 3 * self.head_dim)
+        q, k_new, v_new = torch.chunk(qkv, 3, dim=-1)  # each [B, T_new, H, Dh]
+
+        if past_key_values is not None:
+            k_past, v_past = past_key_values
+            k = torch.cat([k_past, k_new], dim=1)  # [B, T_all, H, Dh]
+            v = torch.cat([v_past, v_new], dim=1)
+        else:
+            k, v = k_new, v_new
+
+        attn_out = flash_attn_func(
+            q, k, v,
+            softmax_scale=1.0 / math.sqrt(self.head_dim),
+            causal=self.causal,
+            window_size=(self.window_size - 1, 0) if self.causal else (-1, -1),
+            alibi_slopes=self.alibi,
+            return_attn_probs=False,
+        )
+        # attn_out: [B, T_new, H, Dh]
+        out = attn_out.contiguous().view(B, T_new, -1)   # [B, T_new, D]
+        y = self.out_proj(out)                           # [B, T_new, D]
+
+        return y, (k, v)
+
+# -----------------------------------------------------------------------------
+# Transformer Blocks and Model
+# -----------------------------------------------------------------------------
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config, att_global: bool = True):
+        super().__init__()
+        if att_global:
+            self.attn = AFTSimple(embed_dim=config.d_model, causal=config.causal)
+        else:
+            self.attn = FlashAttention(
+                embed_dim=config.d_model,
+                num_heads=config.num_attention_heads,
+                window_size=config.window_size,
+                causal=config.causal,
+                qkv_bias=True,
+            )
+        self.ff = nn.Sequential(
+            FeedForward(config.d_model),
+            FeedForward(config.d_model),
+        )
+        self.ln1 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+        self.ln2 = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+        self.drop = nn.Dropout(config.dropout)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
+    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
+        # Attention + residual
+        residual = x
+        x = self.ln1(x)
+        y, present = self.attn(x, past_key_values=past_key_values)
+        x = self.drop(y) + residual
+
+        # Feed-forward + residual
+        residual = x
+        x = self.ln2(x)
+        x = self.ff(x)
+        x = self.drop(x) + residual
+
+        return x, present
+
+class ResidualBlocks(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        blocks = []
+        for i in range(config.num_hidden_layers):
+            # alternate local/global: every 3rd layer global
+            att_global = ((i + 1) % 3 == 0)
+            blocks.append(TransformerBlock(config, att_global=att_global))
+        self.layers = nn.ModuleList(blocks)
+        self.final_ln = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor,torch.Tensor], ...]] = None
+    ) -> Tuple[torch.Tensor, Tuple[Tuple[torch.Tensor,torch.Tensor], ...]]:
+        new_past = []
+        for i, layer in enumerate(self.layers):
+            pkv = None if past_key_values is None else past_key_values[i]
+            x, present = layer(x, past_key_values=pkv)
+            new_past.append(present)
+        x = self.final_ln(x)
+        return x, tuple(new_past)
+
+# -----------------------------------------------------------------------------
+# Configuration and Model
+# -----------------------------------------------------------------------------
+
+class TyneRoxConfig(PretrainedConfig):
+    model_type = "tynerox"
+
+    def __init__(
+        self,
+        vocab_size: int = 30522,
+        context_length: int = 2048,
+        d_model: int = 1024,
+        num_heads: int = 16,
+        window_size: int = 512,
+        num_hidden_layers: int = 12,
+        causal: bool = True,
+        dropout: float = 0.1,
+        layer_norm_eps: float = 1e-5,
+        tie_word_embeddings: bool = False,
+        pad_token_id:int = 0,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = context_length
+        self.d_model = d_model
+        self.num_attention_heads = num_heads
+        self.window_size = window_size
+        self.num_hidden_layers = num_hidden_layers
+        self.causal = causal
+        self.dropout = dropout
+        self.layer_norm_eps = layer_norm_eps
+        self.tie_word_embeddings = tie_word_embeddings
+        self.pad_token_id = pad_token_id
+
+class TyneRoxModel(PreTrainedModel, GenerationMixin):
+    config_class = TyneRoxConfig
+
+    def __init__(self, config: TyneRoxConfig):
+        super().__init__(config)
+        self.embed = PositionalEmbedding(
+            config.vocab_size,
+            config.d_model,
+            config.max_position_embeddings,
+            dropout=config.dropout,
+            encoding_type='rotary'
+        )
+        self.transformer = ResidualBlocks(config)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed.token_embedding
+
+    def set_input_embeddings(self, value):
+        self.embed.token_embedding = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, value):
+        self.lm_head = value
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: bool = True,
+        return_dict: bool = True,
+        **kwargs
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        # 1) Embeddings
+        x = self.embed(input_ids)  # [B, T, D]
+
+        # 2) Transformer blocks with KV-cache
+        x, new_past = self.transformer(x, past_key_values=past_key_values)
+
+        # 3) Project to vocabulary logits
+        logits = self.lm_head(x)  # [B, T, V]
+
+        # 4) Compute loss if labels provided
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        # 5) Return standardized output
+        if not return_dict:
+            output = (logits, new_past) if use_cache else (logits,)
+            return ((loss,) + output) if loss is not None else output
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=new_past if use_cache else None,
+        )
+
+    def _reorder_cache(
+            self,
+            past_key_values: Tuple[Tuple[torch.Tensor, torch.Tensor], ...],
+            beam_idx: torch.Tensor
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
+        reordered = []
+        for k, v in past_key_values:
+            # ambos têm batch dim = dim 0
+            reordered.append((k.index_select(0, beam_idx),
+                              v.index_select(0, beam_idx)))
+        return tuple(reordered)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        **kwargs
+    ) -> dict:
+        # at generation time, only feed in the last token
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:].contiguous()
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+        }

src/visualizations/sample.html

@@ -0,0 +1,723 @@
+<html><head><meta charset='utf-8'></head><body><h4><b>prompt:</b> Olá, como vai você? </h4><p><b>Resposta:</b> Olá, como vai você? .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> Quando a manhã chegou, Iracema ainda estava ali, debruçada, como uma borboleta que </h4><p><b>Resposta:</b> Quando a manhã chegou, Iracema ainda estava ali, debruçada, como uma borboleta que .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> Não, respondeu; na verdade, estou com medo </h4><p><b>Resposta:</b> Não, respondeu; na verdade, estou com medo .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> O resultado representa uma desaceleração </h4><p><b>Resposta:</b> O resultado representa uma desaceleração .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> No vídeo, é possível ver </h4><p><b>Resposta:</b> No vídeo, é possível ver .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> Essa receita de torta de frango </h4><p><b>Resposta:</b> Essa receita de torta de frango .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> Durante o primeiro mandato </h4><p><b>Resposta:</b> Durante o primeiro mandato .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p><h4><b>prompt:</b> Os donos de cães </h4><p><b>Resposta:</b> Os donos de cães .
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+</p></body></html>

tokenizer_config.json

@@ -0,0 +1,91 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|user_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|user_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|assistant_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<|assistant_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<|think_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|think_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<|command_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<|command_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|endoftext|>",
+  "tokenizer_class": "PreTrainedTokenizer",
+  "unk_token": "[UNK]"
+}