Update README.md

4bba615 verified 3 days ago

13.2 kB

language:
  - ru
  - en
  - zh
  - de
  - uk
tags:
  - causal-lm
  - pytorch
  - safetensors
  - timegpt
  - 0.5b
  - text-generation
license: apache-2.0
pipeline_tag: text-generation
inference: false
model_name: TimeGPT-Base

⏳ TimeGPT-Base (0.5B)

TimeGPT-Base — это мультиязычная авторегрессионная языковая модель (Causal LM) класса 0.5B, разработанная и обученная полностью с нуля (from scratch) без использования сторонних весов или дистилляции.

Модель создана независимой командой LongTime Devs и обучена на одном потребительском ускорителе NVIDIA GeForce RTX 5060 Ti 16GB. Основной язык обучения — русский, с интеграцией английского, китайского, немецкого и украинского корпусов.

⚠️ Важное замечание: это базовая фундаментальная модель (Base). Она обучена задаче продолжения текста (next-token prediction / text completion) и не проходила стадию выравнивания инструкций.

Диалоговая инструктивная версия — TimeGPT-Instruct (0.5B) — в настоящее время находится на финальной стадии дообучения и в скором времени будет опубликована в репозитории команды LongTime Devs.

🏛 Архитектурные особенности

TimeGPT построена на современной архитектуре трансформера (уровня Llama 3 и Qwen 2.5), оптимизированной для высокой скорости инференса и экономии видеопамяти:

Grouped-Query Attention (GQA): 20 голов Query и 4 головы Key/Value (соотношение 5:1), что кратно снижает размер KV-кеша при работе с контекстом.
SwiGLU Activation: нелинейная функция активации в блоках MLP.
RMSNorm: быстрая среднеквадратичная нормализация вместо классического LayerNorm.
Rotary Position Embeddings (RoPE): векторное позиционное кодирование с базовой частотой $\theta = 500,000$ для стабильности контекста.
Weight Tying: совмещение матриц эмбеддингов токенов и выходной языковой головы для компактности и регуляризации весов.
FlashAttention / SDPA: поддержка нативного аппаратного ускорения Scaled Dot-Product Attention в режиме bfloat16.

Спецификация модели

Параметр	Значение
Класс параметров	0.5B
Размер файла весов	1.02 GB (`model.safetensors`)
Размерность скрытого слоя (`dim`)	1280
Количество слоев (`n_layers`)	24
Голов внимания (`n_heads`)	20
Голов Key/Value (`n_kv_heads`)	4 (GQA)
Размерность промежуточного слоя MLP	3456 (SwiGLU)
Контекстное окно (`max_seq_len`)	1024 токена
Размер словаря (`vocab_size`)	32 000
Формат весов	SafeTensors (`bfloat16`)

🔤 Токенизатор

Модель использует собственный мультиязычный Byte-Level BPE токенизатор:

Размер словаря: 32 000 токенов.
Byte-fallback: байтовое резервирование исключает появление токенов неизвестных символов (<|unk|>) при обработке редких знаков, иероглифов и эмодзи.
Служебные токены стандарта ChatML:
- <|pad|>
- <|unk|>
- <|im_start|>
- <|im_end|>

📚 Данные и процесс предобучения

Модель обучалась фундаментальному чтению в непрерывном мультиязычном потоке данных:

Русский язык (основной массив, ~55%): русскоязычный срез Википедии + новостные корпуса (IlyaGusev/gazeta).
Английский язык (~20%): статьи англоязычной Википедии.
Китайский язык (~10%): срез китайской Википедии.
Немецкий язык (~8%): срез немецкой Википедии.
Украинский язык (~7%): статьи украинской Википедии.

Детали тренировочного цикла:

Оборудование: 1× NVIDIA GeForce RTX 5060 Ti 16GB.
Шаги обучения: 50 000 шагов (эффективный размер батча = 16 последовательностей по 1024 токена).
Суммарный объем прочитанных данных: ~820 млн токенов.
Время непрерывного обучения: ~27 часов.
Оптимизатор: AdamW ($\beta_1 = 0.9$, $\beta_2 = 0.95$, weight decay = 0.01).
График Learning Rate: Cosine Annealing с плавным разогревом (Warmup: 1500 шагов, Peak LR: 1.8e-4).
Динамика Loss: стартовый лосс 10.64 $\rightarrow$ финальный лосс 3.02–3.36 (перплексия снизилась с бесконечности до ~20).

💻 Быстрый запуск (Inference)

Установите необходимые зависимости:

pip install torch tokenizers safetensors huggingface_hub

Скрипт автоматически скачивает веса (model.safetensors — 1.02 GB) и токенизатор напрямую из репозитория longtimedevs/TimeGPT-Base:

# WARNING: DO NOT USE THIS SCRIPT IN PRODUCTION ENVIRONMENTS! // ВНИМАНИЕ: НЕ ИСПОЛЬЗУЙТЕ ДАННЫЙ КОД В ПРОДАКШЕНЕ!
# FOR RESEARCH, TESTING AND EXPERIMENTAL PURPOSES ONLY! // ТОЛЬКО ДЛЯ ИССЛЕДОВАТЕЛЬСКИХ ЦЕЛЕЙ, ТЕСТИРОВАНИЯ И ЭКСПЕРИМЕНТОВ!

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from tokenizers import Tokenizer
from safetensors.torch import load_file
from huggingface_hub import hf_hub_download

# --- 1. Автоматическое скачивание с Hugging Face ---
REPO_ID = "longtimedevs/TimeGPT-Base"
print(f"[*] Проверка и скачивание файлов из {REPO_ID}...")
weights_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
tokenizer_path = hf_hub_download(repo_id=REPO_ID, filename="tokenizer.json")

# --- 2. Конфигурация модели ---
@dataclass
class TimeGPTConfig:
    vocab_size: int = 32000
    dim: int = 1280
    n_layers: int = 24
    n_heads: int = 20
    n_kv_heads: int = 4
    multiple_of: int = 256
    max_seq_len: int = 1024
    norm_eps: float = 1e-5
    rope_theta: float = 500000.0

# --- 3. Архитектура TimeGPT ---
class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

def precompute_rope_sincos(dim: int, max_seq_len: int, theta: float = 500000.0):
    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
    t = torch.arange(max_seq_len, dtype=torch.float32)
    freqs = torch.outer(t, freqs)
    freqs = torch.cat([freqs, freqs], dim=-1)
    return torch.cos(freqs), torch.sin(freqs)

def rotate_half(x: torch.Tensor):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rope(xq, xk, cos, sin):
    cos = cos[:xq.shape[2]].unsqueeze(0).unsqueeze(1)
    sin = sin[:xq.shape[2]].unsqueeze(0).unsqueeze(1)
    return (xq * cos) + (rotate_half(xq) * sin), (xk * cos) + (rotate_half(xk) * sin)

class Attention(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.n_heads, self.n_kv_heads = cfg.n_heads, cfg.n_kv_heads
        self.head_dim = cfg.dim // cfg.n_heads
        self.n_rep = self.n_heads // self.n_kv_heads
        self.wq = nn.Linear(cfg.dim, cfg.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(cfg.dim, cfg.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(cfg.dim, cfg.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(cfg.n_heads * self.head_dim, cfg.dim, bias=False)

    def forward(self, x, cos, sin):
        B, S, _ = x.shape
        xq = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
        xk = self.wk(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)
        xv = self.wv(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)
        xq, xk = apply_rope(xq, xk, cos, sin)
        if self.n_rep > 1:
            xk = xk.repeat_interleave(self.n_rep, dim=1)
            xv = xv.repeat_interleave(self.n_rep, dim=1)
        out = F.scaled_dot_product_attention(xq, xk, xv, is_causal=True)
        return self.wo(out.transpose(1, 2).contiguous().view(B, S, -1))

class SwiGLUMLP(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        hidden_dim = int(2 * (4 * cfg.dim) / 3)
        hidden_dim = cfg.multiple_of * ((hidden_dim + cfg.multiple_of - 1) // cfg.multiple_of)
        self.w1 = nn.Linear(cfg.dim, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, cfg.dim, bias=False)
        self.w3 = nn.Linear(cfg.dim, hidden_dim, bias=False)

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class TransformerBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.attn_norm, self.attn = RMSNorm(cfg.dim, cfg.norm_eps), Attention(cfg)
        self.ffn_norm, self.ffn = RMSNorm(cfg.dim, cfg.norm_eps), SwiGLUMLP(cfg)

    def forward(self, x, cos, sin):
        return x + self.ffn(self.ffn_norm(x + self.attn(self.attn_norm(x), cos, sin)))

class TimeGPT(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.cfg = cfg
        self.tok_embeddings = nn.Embedding(cfg.vocab_size, cfg.dim)
        self.layers = nn.ModuleList([TransformerBlock(cfg) for _ in range(cfg.n_layers)])
        self.norm = RMSNorm(cfg.dim, cfg.norm_eps)
        self.output = nn.Linear(cfg.dim, cfg.vocab_size, bias=False)
        self.output.weight = self.tok_embeddings.weight
        cos, sin = precompute_rope_sincos(cfg.dim // cfg.n_heads, cfg.max_seq_len, cfg.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=False)
        self.register_buffer("rope_sin", sin, persistent=False)

    def forward(self, idx):
        B, S = idx.shape
        h = self.tok_embeddings(idx)
        cos, sin = self.rope_cos[:S], self.rope_sin[:S]
        for layer in self.layers:
            h = layer(h, cos, sin)
        return self.output(self.norm(h))

# --- 4. Загрузка весов и инференс со стримингом ---
device = "cuda" if torch.cuda.is_available() else "cpu"
cfg = TimeGPTConfig()
model = TimeGPT(cfg).to(device, dtype=torch.bfloat16)

# Чтение скачанных весов SafeTensors
weights = load_file(weights_path, device=device)
model.load_state_dict(weights, strict=False)
model.eval()

tokenizer = Tokenizer.from_file(tokenizer_path)

prompt = "Искусственный интеллект — это"
input_ids = torch.tensor([tokenizer.encode(prompt).ids], device=device)

print(prompt, end="", flush=True)
with torch.inference_mode():
    for _ in range(80):
        logits = model(input_ids[:, -cfg.max_seq_len:])
        next_token = torch.multinomial(F.softmax(logits[0, -1, :] / 0.7, dim=-1), num_samples=1)
        if next_token.item() == tokenizer.token_to_id("<|im_end|>"):
            break
        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)
        print(tokenizer.decode([next_token.item()]), end="", flush=True)
print()

👥 Разработчики

Команда: LongTime Devs
Архитектура: TimeGPT
Лицензия: Apache 2.0 (открыто для свободного использования и исследований)