Python/AI/train_improved.py - main - Demek/DSam621

Python/AI/train_improved.py
379 строк15 KB
d_e_m_e_k
добавил учебную ai для обучения
27 июл 2026, 23:34
27 июл 2026, 23:342dfa071
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import json
import os
from collections import Counter
import math
from pathlib import Path

# ============================================
# 1. Улучшенный токенизатор для YAML
# ============================================

class YAMLTokenizer:
    def __init__(self):
        # Специальные токены
        self.pad_token = "<PAD>"
        self.unk_token = "<UNK>"
        self.bos_token = "<BOS>"  # начало генерации
        self.eos_token = "<EOS>"  # конец генерации

        self.stoi = {
            self.pad_token: 0,
            self.unk_token: 1,
            self.bos_token: 2,
            self.eos_token: 3
        }
        self.itos = {v: k for k, v in self.stoi.items()}

        self.max_len = 128

    def tokenize(self, text):
        """
        Разбивает YAML-подобный текст на токены:
        - ключи: "spec:", "replicas:", "containerPort:"
        - значения: "3", "nginx:latest"
        - разделители: ":", "-", "#", новые строки
        - пробелы игнорируются
        """
        # Разбиваем на строки, затем по символам
        tokens = []
        for line in text.split("\n"):
            # Заменяем специальные символы на отдельные токены
            line = line.replace(":", " : ").replace("-", " - ").replace("#", " # ")
            # Добавляем пробелы вокруг индентации
            line = line.replace("  ", " __INDENT__ ")
            # Делим на слова
            words = line.strip().split()
            tokens.extend(words)
            if line.strip():
                tokens.append("<NL>")  # end-of-line marker

        return tokens

    def build_vocab(self, texts, min_freq=1):
        """Строит словарь по корпусу"""
        freq = Counter()
        for text in texts:
            for token in self.tokenize(text):
                freq[token] += 1

        for token, count in freq.items():
            if count >= min_freq and token not in self.stoi:
                idx = len(self.stoi)
                self.stoi[token] = idx
                self.itos[idx] = token

        print(f"✅ Словарь: {len(self.stoi)} слов")
        return len(self.stoi)

    def encode(self, text):
        """Кодирует текст → ID"""
        tokens = self.tokenize(text)
        ids = [self.stoi.get(tok, self.stoi[self.unk_token]) for tok in tokens]
        return ids

    def decode(self, ids):
        """Декодирует ID → текст"""
        tokens = [self.itos.get(i, self.unk_token) for i in ids]
        # Собираем обратно
        text = " ".join(tokens)
        # Восстанавливаем структуру YAML
        text = text.replace(" __INDENT__ ", "  ").replace(" : ", ":").replace(" - ", "-").replace(" # ", "#")
        text = text.replace(" <NL> ", "\n").replace("<NL>", "\n")
        return text.strip()

    def save(self, path):
        os.makedirs(os.path.dirname(path) if os.path.dirname(path) else ".", exist_ok=True)
        torch.save({"stoi": self.stoi, "itos": self.itos}, path)
        print(f"💾 Словарь сохранён в {path}")

    @classmethod
    def load(cls, path):
        data = torch.load(path)
        tokenizer = cls()
        tokenizer.stoi = data["stoi"]
        tokenizer.itos = data["itos"]
        return tokenizer


# ============================================
# 2. Датасет с правильной обработкой токенов
# ============================================

class K8sDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_len=128):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.data = []

        if not os.path.exists(jsonl_path):
            raise FileNotFoundError(f"❌ Файл {jsonl_path} не найден. Создайте dataset.jsonl.")

        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for i, line in enumerate(f):
                try:
                    item = json.loads(line.strip())
                    self.data.append(item)
                except json.JSONDecodeError as e:
                    print(f"⚠️ Пропущена строка {i}: {e}")

        print(f"✅ Загружено {len(self.data)} примеров")

        if len(self.data) == 0:
            raise ValueError("❌ Датасет пуст!")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        item = self.data[idx]
        instruction = item.get('instruction', '')
        response = item.get('response', '')

        # Формируемprompt-response format
        text = f"{self.tokenizer.bos_token} Инструкция: {instruction}\nОтвет: {response} {self.tokenizer.eos_token}"

        ids = self.tokenizer.encode(text)
        ids = ids[:self.max_len - 2]  # оставляем место для BOS и EOS
        ids = [self.tokenizer.stoi[self.tokenizer.bos_token]] + ids + [self.tokenizer.stoi[self.tokenizer.eos_token]]

        # Дополняем padding'ом справа
        if len(ids) < self.max_len:
            ids = ids + [self.tokenizer.stoi[self.tokenizer.pad_token]] * (self.max_len - len(ids))

        # Создаём x и y: x = без последнего, y = без первого (shifted right)
        x = ids[:-1]
        y = ids[1:]

        return {
            'input_ids': torch.tensor(x, dtype=torch.long),
            'labels': torch.tensor(y, dtype=torch.long)
        }


# ============================================
# 3. Модель: улучшенная LSTM
# ============================================

class MyAIKuberMini(nn.Module):
    def __init__(self, vocab_size=1000, embed_dim=256, hidden_dim=256, num_layers=2, dropout=0.2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(
            input_size=embed_dim,
            hidden_size=hidden_dim,
            num_layers=num_layers,
            batch_first=True,
            bidirectional=False,
            dropout=dropout if num_layers > 1 else 0.0
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, hidden=None):
        x = self.embedding(x)
        x = self.dropout(x)
        lstm_out, hidden = self.lstm(x, hidden)
        logits = self.fc(self.dropout(lstm_out))
        return logits, hidden

    def generate(self, tokenizer, prompt_ids, max_new_tokens=100, temperature=0.8, top_k=50, top_p=0.9, device='cpu'):
        self.eval()
        with torch.no_grad():
            prompt_ids = torch.tensor([prompt_ids], dtype=torch.long).to(device)
            generated = prompt_ids[0].tolist()

            for _ in range(max_new_tokens):
                logits, _ = self(prompt_ids)
                next_logits = logits[:, -1, :] / temperature

                # top-k filtering
                if top_k > 0:
                    indices_to_remove = next_logits < torch.topk(next_logits, top_k)[0][..., -1, None]
                    next_logits[indices_to_remove] = float('-inf')

                # top-p (nucleus) filtering
                if top_p < 1.0:
                    sorted_logits, sorted_indices = torch.sort(next_logits, descending=True)
                    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
                    sorted_indices_to_remove = cumulative_probs > top_p
                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
                    sorted_indices_to_remove[..., 0] = 0
                    indices_to_remove = sorted_indices_to_remove.scatter(-1, sorted_indices, sorted_indices_to_remove)
                    next_logits[indices_to_remove] = float('-inf')

                probs = torch.softmax(next_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1).item()

                generated.append(next_token)
                prompt_ids = torch.tensor([[next_token]], dtype=torch.long).to(device)

                if next_token == tokenizer.stoi[tokenizer.eos_token]:
                    break

            return generated[len(prompt_ids[0]):]  # убираем prompt из результата


# ============================================
# 4. Обучение и валидация
# ============================================

def train_epoch(model, dataloader, optimizer, criterion, device, accumulation_steps=4):
    model.train()
    total_loss = 0
    optimizer.zero_grad()

    for i, batch in enumerate(dataloader):
        input_ids = batch['input_ids'].to(device)
        labels = batch['labels'].to(device)

        logits, _ = model(input_ids)
        loss = criterion(logits.view(-1, logits.size(-1)), labels.view(-1))
        loss = loss / accumulation_steps
        loss.backward()

        if (i + 1) % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            optimizer.step()
            optimizer.zero_grad()

        total_loss += loss.item() * accumulation_steps

    return total_loss / len(dataloader)


def validate(model, dataloader, criterion, device):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for batch in dataloader:
            input_ids = batch['input_ids'].to(device)
            labels = batch['labels'].to(device)
            logits, _ = model(input_ids)
            loss = criterion(logits.view(-1, logits.size(-1)), labels.view(-1))
            total_loss += loss.item()
    return total_loss / len(dataloader)


# ============================================
# 5. Генерация YAML
# ============================================

def generate_yaml(model, tokenizer, instruction, max_new_tokens=100, temperature=0.8, device='cpu'):
    prompt = f"{tokenizer.bos_token} Инструкция: {instruction}\nОтвет:"
    prompt_ids = tokenizer.encode(prompt)

    generated_ids = model.generate(tokenizer, prompt_ids, max_new_tokens, temperature, top_k=50, top_p=0.9, device=device)

    full_text = tokenizer.decode(prompt_ids + generated_ids)

    if "Ответ:" in full_text:
        yaml_part = full_text.split("Ответ:")[1].strip()
        yaml_part = yaml_part.replace(tokenizer.eos_token, "").strip()
    else:
        yaml_part = full_text.replace(tokenizer.bos_token, "").strip()

    return yaml_part


# ============================================
# 6. Главная функция
# ============================================

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"🚀 Используем устройство: {device}")

    # Создаём примеры (если файла нет)
    dataset_path = "data/expanded_dataset.jsonl"
    if not os.path.exists(dataset_path):
        examples = [
            {
                "instruction": "Создай Deployment для nginx с 3 репликами",
                "response": "apiVersion: apps/v1\nkind: Deployment\nmetadata:\n  name: nginx\nspec:\n  replicas: 3\n  selector:\n    matchLabels:\n      app: nginx\n  template:\n    metadata:\n      labels:\n        app: nginx\n    spec:\n      containers:\n      - name: nginx\n        image: nginx:latest\n        ports:\n        - containerPort: 80"
            },
            {
                "instruction": "Создай Service для nginx",
                "response": "apiVersion: v1\nkind: Service\nmetadata:\n  name: nginx-svc\nspec:\n  selector:\n    app: nginx\n  ports:\n  - port: 80\n    targetPort: 80\n  type: ClusterIP"
            },
        ]
        os.makedirs("data", exist_ok=True)
        with open(dataset_path, 'w', encoding='utf-8') as f:
            for ex in examples:
                f.write(json.dumps(ex, ensure_ascii=False) + '\n')
        print(f"✅ dataset.jsonl создан ({len(examples)} примеров)")

    # Инициализация токенизатора
    tokenizer_path = "data/tokenizer.pth"
    if os.path.exists(tokenizer_path):
        print(f"📂 Загрузка токенизатора из {tokenizer_path}")
        tokenizer = YAMLTokenizer.load(tokenizer_path)
    else:
        tokenizer = YAMLTokenizer()

        texts = []
        with open(dataset_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                texts.append(f"{item.get('instruction', '')} {item.get('response', '')}")

        vocab_size = tokenizer.build_vocab(texts, min_freq=1)
        tokenizer.save(tokenizer_path)

    # Данные и DataLoader
    dataset = K8sDataset(dataset_path, tokenizer, max_len=256)
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])

    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=0, drop_last=True)
    val_loader = DataLoader(val_dataset, batch_size=2, shuffle=False, num_workers=0, drop_last=False)

    # Модель
    model = MyAIKuberMini(
        vocab_size=len(tokenizer.stoi),
        embed_dim=256,
        hidden_dim=256,
        num_layers=2
    ).to(device)

    params = sum(p.numel() for p in model.parameters())
    print(f"🔢 Параметров: {params:,}")

    # Обучение
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi[tokenizer.pad_token])
    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

    best_val_loss = float('inf')
    os.makedirs("checkpoints", exist_ok=True)

    print("\n🔄 Обучение...")
    for epoch in range(100):
        train_loss = train_epoch(model, train_loader, optimizer, criterion, device)
        val_loss = validate(model, val_loader, criterion, device)
        scheduler.step()

        print(f"Epoch {epoch+1:3d} | Train Loss: {train_loss:.4f} | Val Loss: {val_loss:.4f}")

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), "checkpoints/best_model.pth")
            print("💾 Сохранена лучшая модель (Val Loss: {val_loss:.4f})")

    # Итоговая генерация
    print("\n🚀 Генерация YAML...")
    instruction = "Создай Deployment для nginx с 3 репликами"
    yaml = generate_yaml(model, tokenizer, instruction, temperature=0.7, device=device)
    print(f"\n📌 Инструкция: {instruction}")
    print(f"📝 YAML:\n{yaml}")

    # Сохранение модели
    torch.save(model.state_dict(), "model_weights_improved.pth")
    print(f"\n✅ Модель сохранена в model_weights_improved.pth")


if __name__ == "__main__":
    main()