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

Python/AI/kuber1b.py
210 строк7 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

class SimpleTokenizer:
    def __init__(self):
        self.stoi = {"<PAD>": 0, "<UNK>": 1}
        self.itos = {0: "<PAD>", 1: "<UNK>"}
    
    def build_vocab(self, texts, min_freq=2):
        from collections import Counter
        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)} слов")
    
    def tokenize(self, text):
        return text.lower().split()
    
    def encode(self, text):
        return [self.stoi.get(tok, self.stoi["<UNK>"]) for tok in self.tokenize(text)]
    
    def decode(self, ids):
        return ' '.join([self.itos.get(i, "<UNK>") for i in ids])

class K8sDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_len=128):
        self.tokenizer = tokenizer
        self.max_len = max_len
        self.data = []
        
        with open(jsonl_path, 'r', encoding='utf-8') as f:
            for line in f:
                item = json.loads(line.strip())
                self.data.append(item)
        
        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', '')
        
        text = f"Инструкция: {instruction}\nОтвет: {response}"
        
        ids = self.tokenizer.encode(text)
        
        if len(ids) > self.max_len:
            ids = ids[:self.max_len]
        else:
            ids = ids + [0] * (self.max_len - len(ids))
        
        return torch.tensor(ids, dtype=torch.long)

class MyAIKuberMini(nn.Module):
    def __init__(self, vocab_size = 220000, embed_dim = 2048, hidden_dim = 2048, num_layers=4):
        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)
        lstm_out, _ = self.lstm(x)
        logits = self.fc(lstm_out)
        return logits

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):
        batch = batch.to(device)
        
        x = batch[:, :-1]  
        y = batch[:, 1:]  
        
        logits = model(x) 
        loss = criterion(logits.reshape(-1, logits.size(-1)), y.reshape(-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
        
        if i % 10 == 0:
            torch.cuda.empty_cache()
    
    return total_loss / len(dataloader)

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.encode(prompt)
        tokens = torch.tensor([ids], dtype=torch.long).to(device)
        
        generated_ids = ids.copy()
        
        for _ in range(max_new_tokens):
            logits = model(tokens) 
            
            next_logits = logits[:, -1, :] / temperature
            probs = torch.softmax(next_logits, dim=-1)
            
            next_token = torch.multinomial(probs, num_samples=1).item()
            
            generated_ids.append(next_token)
            tokens = torch.tensor([generated_ids], dtype=torch.long).to(device)
            
            if next_token == tokenizer.stoi.get("<EOS>", 0):
                break
        
        full_text = tokenizer.decode(generated_ids)
        
        if "Ответ:" in full_text:
            yaml_part = full_text.split("Ответ:")[1].strip()
        else:
            yaml_part = full_text
        
        return yaml_part

def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Используем: {device}")
    
    if not os.path.exists('dataset.jsonl'):
        print("примеры")
        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", "path": "nginx.yaml"},
            {"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", "path": "service.yaml"},
        ]
        with open('dataset.jsonl', '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.jsonl', 'r', encoding='utf-8') as f:
        for line in f:
            item = json.loads(line.strip())
            texts.append(f"{item.get('instruction', '')} {item.get('response', '')}")
    
    tokenizer.build_vocab(texts)
    vocab_size = len(tokenizer.stoi)
    print(f"Размер словаря: {vocab_size}")
    
    model = MyAIKuberMini(
        vocab_size=vocab_size,
        embed_dim=256,
        hidden_dim=256,
        num_layers=2
    ).half().to(device)
    
    params = sum(p.numel() for p in model.parameters())
    print(f"Параметров: {params:,}")
    
    dataset = K8sDataset('dataset.jsonl', tokenizer, max_len=64)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
    
    criterion = nn.CrossEntropyLoss(ignore_index=0)
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-4)
    
    print("\nобучение")
    for epoch in range(100):
        loss = train_epoch(model, dataloader, optimizer, criterion, device)
        print(f"Epoch {epoch+1}, Loss: {loss:.4f}")
    
    print("\nитог")
    instruction = "Создай Service для nginx"
    yaml = generate_yaml(model, tokenizer, instruction, device=device)
    print(f"Инструкция: {instruction}")
    print(f"YAML:\n{yaml}")
    
    torch.save(model.state_dict(), 'model_weights.pth')
    print("\nМодель сохранена в model_weights.pth")

if __name__ == "__main__":
    main()