Python/AI/kuber1b-fix.py - main - Demek/DSam621

Python/AI/kuber1b-fix.py
285 строк11 KB
d_e_m_e_k
ДЗ10
02 авг 2026, 20:39
02 авг 2026, 20:398671c2f
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
import json
import os

# ==========================================
# 1. Токенизатор с BOS и EOS
# ==========================================
class SimpleTokenizer:
    def __init__(self):
        # 0: PAD, 1: UNK, 2: BOS, 3: EOS
        self.stoi = {"<PAD>": 0, "<UNK>": 1, "<BOS>": 2, "<EOS>": 3}
        self.itos = {v: k for k, v in self.stoi.items()}
    
    def build_vocab(self, texts, min_freq=1):
        from collections import Counter
        freq = Counter()
        for text in texts:
            # Простой split по пробелам
            tokens = text.split()
            for token in tokens:
                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)} слов")
    
    def tokenize(self, text):
        # Добавляем BOS в начало и EOS в конец
        tokens = ["<BOS>"] + text.split() + ["<EOS>"]
        return tokens
    
    def encode(self, text, max_len=256):
        tokens = self.tokenize(text)
        ids = [self.stoi.get(tok, self.stoi["<UNK>"]) for tok in tokens]
        
        # Обрезка или PAD
        if len(ids) > max_len:
            ids = ids[:max_len]
        else:
            ids = ids + [self.stoi["<PAD>"]] * (max_len - len(ids))
        
        return ids
    
    def decode(self, ids):
        # Декодируем, игнорируя PAD
        tokens = [self.itos.get(i, "<UNK>") for i in ids if i != self.stoi["<PAD>"]]
        return ' '.join(tokens)

# ==========================================
# 2. Датасет (Правильное сдвигание для LM)
# ==========================================
class K8sDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_len=256):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.data = []
        
        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    try:
                        item = json.loads(line.strip())
                        self.data.append(item)
                    except:
                        continue
        
        print(f"Загружено {len(self.data)} примеров")
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        item = self.data[idx]
        instruction = item.get('instruction', '')
        response = item.get('response', '')
        
        # Формируем полный текст для обучения
        # Модель будет учиться предсказывать: "Инструкция: ... \n Ответ: ... <EOS>"
        full_text = f"Инструкция: {instruction}\nОтвет: {response}"
        
        # Получаем ID всего текста
        ids = self.tokenizer.encode(full_text, self.max_len)
        
        # Для Language Modeling:
        # x (input) : все токены кроме последнего
        # y (target): все токены кроме первого (сдвиг вправо)
        x = torch.tensor(ids[:-1], dtype=torch.long)
        y = torch.tensor(ids[1:], dtype=torch.long)
        
        return x, y

# ==========================================
# 3. Модель (ИСПРАВЛЕННОЕ)
# ==========================================
class MyAIKuberMini(nn.Module):
    def __init__(self, vocab_size, embed_dim=128, hidden_dim=128, num_layers=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=0.2
        )
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        x = self.embedding(x)
        x = self.dropout(x)
        # Возвращаем outputs и скрытое состояние (h_n, c_n)
        lstm_out, (hidden, cell) = self.lstm(x)
        logits = self.fc(self.dropout(lstm_out))
        # ✅ Важно: возвращаем кортеж из двух элементов
        return logits, (hidden, cell)

# ==========================================
# 4. Обучение (ИСПРАВЛЕННОЕ)
# ==========================================
def train_epoch(model, dataloader, optimizer, criterion, device):
    model.train()
    total_loss = 0
    batch_count = 0
    
    for x_batch, y_batch in dataloader:
        x_batch = x_batch.to(device)
        y_batch = y_batch.to(device)
        
        # ✅ Распаковываем результат модели: получаем только logits для loss
        logits, _ = model(x_batch) 
        
        # Перестраиваем для CrossEntropy
        # logits: (batch, seq, vocab) -> (batch*seq, vocab)
        loss = criterion(
            logits.reshape(-1, logits.size(-1)), 
            y_batch.reshape(-1)
        )
        
        loss.backward()
        
        # Градиентный клиппинг - ОЧЕНЬ ВАЖНО для LSTM
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        
        optimizer.step()
        optimizer.zero_grad()
        
        total_loss += loss.item()
        batch_count += 1

    return total_loss / batch_count

# ==========================================
# 5. Генерация (ИСПРАВЛЕННАЯ)
# ==========================================
def generate_yaml(model, tokenizer, instruction, max_new_tokens=100, temperature=0.8, device='cpu'):
    model.eval()
    with torch.no_grad():
        # На вход подаем только инструкцию
        prompt = f"Инструкция: {instruction}\n"
        
        # Кодирование промпта
        ids = tokenizer.tokenize(prompt) 
        token_ids = [tokenizer.stoi[t] for t in ids]
        
        # Передаем промпт в модель
        input_tensor = torch.tensor([token_ids], dtype=torch.long).to(device)
        # ✅ Распаковываем кортеж: берем только logits
        logits, _ = model(input_tensor)
        
        # Берем логиты последнего токена промпта
        next_token_logits = logits[0, -1, :] / temperature
        probs = torch.softmax(next_token_logits, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1).item()
        
        generated_ids = token_ids + [next_token]
        
        # Генерируем дальше
        for _ in range(max_new_tokens):
            curr_input = torch.tensor([generated_ids[-100:]], dtype=torch.long).to(device)
            # ✅ Распаковываем кортеж внутри цикла генерации
            curr_logits, _ = model(curr_input)
            
            next_token_logits = curr_logits[0, -1, :] / temperature
            probs = torch.softmax(next_token_logits, dim=-1)
            
            next_token = torch.multinomial(probs, num_samples=1).item()
            
            generated_ids.append(next_token)
            
            if next_token == tokenizer.stoi["<EOS>"]:
                break
        
        # Декодируем результат
        full_text = tokenizer.decode(generated_ids)
        
        if "Ответ:" in full_text:
            yaml_part = full_text.split("Ответ:")[1].strip()
        else:
            yaml_part = full_text
            
        yaml_part = yaml_part.replace("<BOS>", "").replace("<EOS>", "")
        
        return yaml_part

# ==========================================
# MAIN
# ==========================================
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Используем: {device}")
    
    dataset_path = 'dataset.jsonl'
    
    if not os.path.exists(dataset_path):
        print("Файл dataset.jsonl не найден. Создаю тестовый...")
        examples = [
            {"instruction": "Создай Deployment для nginx", "response": "apiVersion: apps/v1\nkind: Deployment\nname: nginx"},
            {"instruction": "Создай Service для nginx", "response": "apiVersion: v1\nkind: Service\nname: nginx-svc"},
        ]
        with open(dataset_path, 'w', encoding='utf-8') as f:
            for ex in examples:
                f.write(json.dumps(ex, ensure_ascii=False) + '\n')
        print("Тестовый dataset.jsonl создан.")

    tokenizer = SimpleTokenizer()
    
    texts = []
    with open(dataset_path, 'r', encoding='utf-8') as f:
        for line in f:
            if line.strip():
                item = json.loads(line.strip())
                full = f"Инструкция: {item['instruction']}\nОтвет: {item['response']}"
                texts.append(full)
    
    tokenizer.build_vocab(texts, min_freq=1)
    vocab_size = len(tokenizer.stoi)
    print(f"Размер словаря: {vocab_size}")
    
    model = MyAIKuberMini(
        vocab_size=vocab_size,
        embed_dim=64,   # Уменьшил для стабильности
        hidden_dim=64,
        num_layers=2
    ).to(device)
    
    dataset = K8sDataset(dataset_path, tokenizer, max_len=256)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
    
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3) # LR 0.002
    criterion = nn.CrossEntropyLoss(ignore_index=tokenizer.stoi["<PAD>"]) # Игнорируем PAD
    
    print("Обучение...")
    for epoch in range(20):
        loss = train_epoch(model, dataloader, optimizer, criterion, device)
        print(f"Epoch {epoch+1}, Loss: {loss:.4f}")
        
        if torch.isnan(torch.tensor(loss)):
            print("Обнаружен NaN! Прекращение.")
            break

    print("Генерация...")
    instruction = "Создай Deployment для postgres"
    yaml = generate_yaml(model, tokenizer, instruction, device=device)
    print(f"Инструкция: {instruction}")
    print(f"YAML:\n{yaml}")
    
    # Сохраняем словарь токенов
    torch.save({
        'stoi': tokenizer.stoi,
        'itos': tokenizer.itos,
    }, 'tokenizer_weights.pth')
    
    torch.save(model.state_dict(), 'model_weights_fixed.pth')
    print("Модель сохранена.")

if __name__ == "__main__":
    main()