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

Python/AI/infer.py
179 строк7 KB
d_e_m_e_k
ДЗ10
02 авг 2026, 20:39
02 авг 2026, 20:398671c2f
import torch
import torch.nn as nn
import json
import os

# ==========================================
# 1. Классы и функции из обучающего скрипта
# ==========================================
class SimpleTokenizer:
    def __init__(self):
        self.stoi = {"<PAD>": 0, "<UNK>": 1, "<BOS>": 2, "<EOS>": 3}
        self.itos = {v: k for k, v in self.stoi.items()}
    
    # Для загрузки из файла нам нужно знать структуру, 
    # но обычно мы сохраняем словари вместе с моделью или восстанавливаем их.
    # В данном примере мы предполагая, что словарь уже встроен в веса или его нужно построить заново из тех же данных.
    # Чтобы не возиться с сохранением словаря отдельно, мы изменим подход:
    # Мы загрузим веса, но для токенизации нам нужен словот.
    # Простой способ: сохранить словарь отдельно при обучении (см. в конце инструкции).
    
    # Для демо используем жестко закодированные методы, но лучше загружать stoi/itos из файла.
    def build_vocab_from_file(self, path):
        texts = []
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    item = json.loads(line.strip())
                    texts.append(f"Инструкция: {item['instruction']}\nОтвет: {item['response']}")
        
        from collections import Counter
        freq = Counter()
        for text in texts:
            tokens = text.split()
            for token in tokens:
                freq[token] += 1
        
        for token, count in freq.items():
            if count >= 1 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):
        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]
        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):
        tokens = [self.itos.get(i, "<UNK>") for i in ids if i != self.stoi["<PAD>"]]
        return ' '.join(tokens)

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)
        lstm_out, (hidden, cell) = self.lstm(x)
        logits = self.fc(self.dropout(lstm_out))
        return logits, (hidden, cell)

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, _ = 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

# ==========================================
# 2. Главная функция
# ==========================================
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Используем устройство: {device}")

    # 1. Загружаем словарь из датасета, чтобы модель могла его понять
    # ВАЖНО: vocab_size должен совпадать с тем, что был при обучении
    tokenizer = SimpleTokenizer()
    # Загружаем сохраненный словарь
    tokenizer_data = torch.load('tokenizer_weights.pth', map_location=device)
    tokenizer.stoi = tokenizer_data['stoi']
    tokenizer.itos = tokenizer_data['itos']
    vocab_size = len(tokenizer.stoi)

    # 2. Инициализируем модель с теми же параметрами, что и при обучении
    model = MyAIKuberMini(
        vocab_size=vocab_size,
        embed_dim=64,
        hidden_dim=64,
        num_layers=2
    ).to(device)

    # 3. Загружаем веса
    weights_path = 'model_weights_fixed.pth'
    if not os.path.exists(weights_path):
        print(f"Ошибка: Файл {weights_path} не найден! Сначала обучите модель.")
        return

    model.load_state_dict(torch.load(weights_path, map_location=device))
    print(f"Модель успешно загружена с файла: {weights_path}")

    # 4. Интерактивный цикл
    print("\n--- Интерактивный режим ---")
    print("Введите инструкцию (или 'exit' для выхода):")
    
    while True:
        instruction = input("\nВаш запрос: ")
        
        if instruction.lower() in ['exit', 'quit', 'выход']:
            print("До свидания!")
            break
        
        if not instruction.strip():
            continue
            
        try:
            yaml_result = generate_yaml(model, tokenizer, instruction, device=device)
            print("\n--- Ответ модели ---")
            print(yaml_result)
            print("--------------------")
        except Exception as e:
            print(f"Ошибка при генерации: {e}")

if __name__ == "__main__":
    main()