app.py

import torch
import torch.nn as nn
import torch.optim as optim
import random
import os
import gradio as gr
from datetime import datetime
import math
import spaces
import warnings
import time

# Подавляем технические предупреждения
warnings.filterwarnings("ignore", message="The PyTorch API of nested tensors is in prototype stage")
warnings.filterwarnings("ignore", category=UserWarning, module="torch.nn.modules.transformer")

# ============ НАСТРОЙКИ ============
DATA_DIR = '/data'
os.makedirs(DATA_DIR, exist_ok=True)
MODEL_BIN = os.path.join(DATA_DIR, 'andrey_final.bin')
MODEL_GGUF = os.path.join(DATA_DIR, 'andrey_final.gguf')

# ============ СЛОВАРЬ ============
WORDS = [
    'привет', 'здравствуй', 'добрый', 'день', 'утро', 'вечер', 'ночь', 'рад', 'видеть',
    'как', 'ты', 'дела', 'жизнь', 'настроение', 'что', 'нового', 'кто', 'где', 'когда', 'почему',
    'хорошо', 'отлично', 'нормально', 'плохо', 'спасибо', 'пожалуйста', 'да', 'нет',
    'я', 'мы', 'он', 'она', 'думаю', 'знаю', 'хочу', 'могу', 'буду', 'люблю', 'работаю', 'учусь',
    'андрей', 'бот', 'помощник', 'нейросеть', 'евгений', 'создатель', 'openrussianai', 'тверь', 'россия',
    'hugging', 'face', 'платформа', 'дом', 'мир', 'люди', 'друг', 'время', 'сегодня', 'завтра',
    'плюс', 'минус', 'умножить', 'разделить', 'равно', 'ноль', 'один', 'два', 'три', 'четыре', 'пять',
    'и', 'а', 'но', 'или', 'в', 'на', 'с', 'из', 'к', 'у', 'не', 'ли', 'же', 'бы', 'очень', 'уже', 'ещё',
    'мой', 'твой', 'наш', 'этот', 'тот', 'какой', 'который', 'весь', 'сам',
    'расскажи', 'покажи', 'объясни', 'помоги', 'скажи', 'шутка', 'смех', 'радость', 'счастье', 'удача'
]

word_to_idx = {w: i+3 for i, w in enumerate(WORDS)}
idx_to_word = {i+3: w for i, w in enumerate(WORDS)}
idx_to_word[0] = '[PAD]'
idx_to_word[1] = '[UNK]'
idx_to_word[2] = '[START]'

vocab_size = len(WORDS) + 3
PAD, UNK, START = 0, 1, 2
MAX_LEN = 25

def tokenize(text):
    return [word_to_idx.get(w, UNK) for w in text.lower().split()]

def detokenize(tokens):
    words = []
    for t in tokens:
        t = int(t)
        if t in [PAD, UNK, START]:
            continue
        w = idx_to_word.get(t, '')
        if w and not w.startswith('['):
            words.append(w)
    return ' '.join(words)

def pad_sequence(seq, max_len=MAX_LEN):
    return (seq + [PAD] * max_len)[:max_len]

# ============ АРХИТЕКТУРА ============
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer('pe', pe.unsqueeze(0))
        
    def forward(self, x):
        return self.dropout(x + self.pe[:, :x.size(1), :])

class AndreyTransformer(nn.Module):
    def __init__(self, vocab_size=vocab_size, d_model=128, nhead=4, num_layers=2, dim_feedforward=256, dropout=0.1):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoder = PositionalEncoding(d_model, dropout)
        self.pos_decoder = PositionalEncoding(d_model, dropout)
        
        self.transformer = nn.Transformer(
            d_model=d_model, nhead=nhead, num_encoder_layers=num_layers,
            num_decoder_layers=num_layers, dim_feedforward=dim_feedforward,
            dropout=dropout, batch_first=True
        )
        self.fc_out = nn.Linear(d_model, vocab_size)
        self.d_model = d_model

    def generate_mask(self, tgt_len, device):
        return torch.triu(torch.ones(tgt_len, tgt_len, device=device), diagonal=1).bool()
    
    def create_pad_mask(self, seq, pad_idx=PAD, device=None):
        if device is None: device = seq.device
        return (seq == pad_idx).to(device)
        
    def forward(self, src, tgt, src_mask=None, tgt_mask=None, 
                src_key_padding_mask=None, tgt_key_padding_mask=None):
        src_emb = self.pos_encoder(self.embedding(src) * math.sqrt(self.d_model))
        tgt_emb = self.pos_decoder(self.embedding(tgt) * math.sqrt(self.d_model))
        output = self.transformer(
            src_emb, tgt_emb, src_mask=src_mask, tgt_mask=tgt_mask,
            src_key_padding_mask=src_key_padding_mask,
            tgt_key_padding_mask=tgt_key_padding_mask
        )
        return self.fc_out(output)
    
    def encode(self, src):
        src_emb = self.pos_encoder(self.embedding(src) * math.sqrt(self.d_model))
        src_key_padding_mask = self.create_pad_mask(src, device=src.device)
        return self.transformer.encoder(src_emb, src_key_padding_mask=src_key_padding_mask)
    
    def decode_step(self, tgt, memory, tgt_mask=None, tgt_key_padding_mask=None):
        tgt_emb = self.pos_decoder(self.embedding(tgt) * math.sqrt(self.d_model))
        output = self.transformer.decoder(
            tgt_emb, memory, tgt_mask=tgt_mask,
            tgt_key_padding_mask=tgt_key_padding_mask
        )
        return self.fc_out(output)

# ============ ДИАЛОГИ ============
DIALOGUES = [
    ("привет", "привет как твои дела"),
    ("здравствуй", "здравствуй рад тебя видеть"),
    ("доброе утро", "доброе утро хорошего дня"),
    ("добрый день", "добрый день чем могу помочь"),
    ("как дела", "у меня всё хорошо а у тебя"),
    ("как ты", "я отлично спасибо что спросил"),
    ("кто ты", "я андрей искусственный интеллект помощник"),
    ("как тебя зовут", "меня зовут андрей"),
    ("где ты живёшь", "я живу на платформе hugging face"),
    ("кто тебя создал", "меня создал евгений из openrussianai"),
    ("что такое openrussianai", "это компания разработчиков из россии"),
    ("ты из россии", "да мой создатель из твери россия"),
    ("сколько будет два плюс два", "два плюс два равно четыре"),
    ("расскажи шутку", "почему программисты путают хэллоуин и рождество"),
    ("мне грустно", "не грусти всё будет хорошо я рядом"),
    ("спасибо", "пожалуйста обращайся если нужна помощь"),
    ("пока", "до свидания хорошего настроения"),
    ("что ты умеешь", "я умею отвечать на вопросы и считать"),
    ("ты умный", "я стараюсь учиться каждый день"),
    ("мы друзья", "конечно я рад нашей дружбе"),
    ("что такое любовь", "любовь это забота и внимание к другому"),
    ("помоги мне", "конечно скажи что именно нужно сделать"),
    ("ты спишь", "нет я всегда на связи и готов помочь"),
    ("ты ешь", "нет мне не нужна обычная еда"),
    ("ты мечтаешь", "я мечтаю стать самым полезным помощником"),
    ("ты злой", "нет я добрый и дружелюбный бот"),
    ("ты хороший", "спасибо я стараюсь быть хорошим"),
    ("ты красивый", "спасибо мой код очень элегантен"),
    ("ты быстрый", "да я обрабатываю информацию мгновенно"),
    ("ты сильный", "моя сила в знаниях и алгоритмах"),
]

FALLBACK_DICT = {q.lower(): a for q, a in DIALOGUES}

def prepare_data():
    X_questions, Y_answers_input, Y_answers_target = [], [], []
    for q, a in DIALOGUES:
        q_toks, a_toks = tokenize(q), tokenize(a)
        if q_toks and a_toks:
            X_questions.append(pad_sequence(q_toks, MAX_LEN))
            a_input = pad_sequence([START] + a_toks, MAX_LEN)
            a_target = pad_sequence(a_toks + [PAD], MAX_LEN)
            Y_answers_input.append(a_input)
            Y_answers_target.append(a_target)
    return (torch.tensor(X_questions, dtype=torch.long),
            torch.tensor(Y_answers_input, dtype=torch.long),
            torch.tensor(Y_answers_target, dtype=torch.long))

# ============ КЛАСС ANDREY AI ============
class AndreyAI:
    def __init__(self, bin_file=MODEL_BIN, gguf_file=MODEL_GGUF):
        self.bin_file = bin_file
        self.gguf_file = gguf_file
        self.memory = {'chat_history': [], 'epochs_trained': 0}
        self.model = AndreyTransformer() 
        
    def load_weights(self):
        if os.path.exists(self.bin_file):
            try:
                data = torch.load(self.bin_file, map_location='cpu')
                self.model.load_state_dict(data['model_state'])
                self.memory = data.get('memory', self.memory)
                print(f"✅ Веса загружены из {self.bin_file}")
                return True
            except Exception as e: 
                print(f"⚠️ Ошибка загрузки: {e}")
        return False

    def save_weights(self):
        os.makedirs(os.path.dirname(self.bin_file), exist_ok=True)
        state = {
            'model_state': self.model.state_dict(),
            'vocab_size': vocab_size, 'd_model': 128, 'nhead': 4,
            'num_layers': 2, 'dim_feedforward': 256,
            'word_to_idx': word_to_idx,
            'idx_to_word': {str(k): v for k, v in idx_to_word.items()},
            'memory': self.memory, 'version': '11.1-CPU-Train-Fixed',
            'created': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }
        torch.save(state, self.bin_file)
        print(f"✅ Сохранено .bin: {os.path.getsize(self.bin_file)/1024:.1f} КБ")

    # ИСПРАВЛЕНИЕ: Явно указываем 'cpu', чтобы обойти блокировку ZeroGPU
    def train(self, epochs=50):
        print(f"🚀 Обучение ({epochs} эпох) на CPU (быстро и без лимитов ZeroGPU)...")
        device = torch.device('cpu') # <-- КЛЮЧЕВОЕ ИЗМЕНЕНИЕ
        self.model.to(device)
        self.model.train()
        
        X_data, Y_input, Y_target = prepare_data()
        criterion = nn.CrossEntropyLoss(ignore_index=PAD)
        optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        
        for epoch in range(epochs):
            total_loss, n_batches = 0, 0
            indices = list(range(len(X_data)))
            random.shuffle(indices)
            
            for idx in indices:
                src = X_data[idx].unsqueeze(0).to(device)
                tgt_in = Y_input[idx].unsqueeze(0).to(device)
                tgt_tar = Y_target[idx].unsqueeze(0).to(device)
                
                tgt_mask = self.model.generate_mask(tgt_in.size(1), device)
                src_pad = self.model.create_pad_mask(src, device=device)
                tgt_pad = self.model.create_pad_mask(tgt_in, device=device)
                
                optimizer.zero_grad()
                output = self.model(src, tgt_in, tgt_mask=tgt_mask, src_key_padding_mask=src_pad, tgt_key_padding_mask=tgt_pad)
                loss = criterion(output.view(-1, vocab_size), tgt_tar.view(-1))
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                optimizer.step()
                
                total_loss += loss.item()
                n_batches += 1
            
            if epoch % 10 == 0: 
                print(f"Эпоха {epoch}/{epochs} | Loss: {total_loss/n_batches:.4f}")
        
        self.memory['epochs_trained'] += epochs
        self.save_weights()
        return f"✅ Обучение завершено! Эпох: {self.memory['epochs_trained']}"

    def get_fallback_answer(self, question):
        q_clean = question.lower().strip()
        if q_clean in FALLBACK_DICT: return FALLBACK_DICT[q_clean]
        for q, a in DIALOGUES:
            if q in q_clean or q_clean in q: return a
        return "Я пока учусь, но стараюсь понимать тебя."

    # Генерация ОСТАЕТСЯ под @spaces.GPU, так как это единственный способ использовать GPU в ZeroGPU
    @spaces.GPU(duration=10)
    def generate_tokens(self, question, history=None, temperature=0.3, max_length=15):
        device = torch.device('cuda')
        self.model.to(device)
        self.model.eval()
        
        q = question.lower().strip()
        ctx_tokens = tokenize(q)
        if not ctx_tokens: 
            self.model.cpu()
            return [self.get_fallback_answer(q)]
            
        if len(ctx_tokens) > MAX_LEN: ctx_tokens = ctx_tokens[-MAX_LEN:]
        
        src = torch.tensor([pad_sequence(ctx_tokens, MAX_LEN)], dtype=torch.long).to(device)
        response_tokens = []
        last_token = -1
        
        try:
            with torch.no_grad():
                memory = self.model.encode(src)
                decoder_input = torch.tensor([[START]], dtype=torch.long).to(device)
                
                for i in range(max_length):
                    tgt_mask = self.model.generate_mask(decoder_input.size(1), device)
                    output = self.model.decode_step(decoder_input, memory, tgt_mask=tgt_mask)
                    logits = output[:, -1, :] / temperature
                    probs = torch.softmax(logits, dim=-1)
                    _, next_token = torch.max(probs, dim=-1)
                    next_token = next_token.item()
                    
                    if next_token in [PAD, UNK, START] or next_token == last_token:
                        break
                        
                    last_token = next_token
                    word = idx_to_word.get(next_token, '')
                    if word and not word.startswith('['):
                        response_tokens.append(word)
                    
                    decoder_input = torch.cat([decoder_input, torch.tensor([[next_token]], dtype=torch.long).to(device)], dim=1)
        except Exception as e:
            print(f"Ошибка генерации: {e}")
        finally:
            self.model.cpu()

        if not response_tokens:
            return [self.get_fallback_answer(q)]
            
        return response_tokens

# ============ GRADIO ИНТЕРФЕЙС ============
def gradio_chat(message, history):
    if not message:
        return "", history or []
    if history is None:
        history = []
    
    history = history + [{"role": "user", "content": message}]
    history.append({"role": "assistant", "content": ""})
    
    tokens = andrey.generate_tokens(message, history)
    
    current_text = ""
    for token in tokens:
        current_text += token + " "
        history[-1]["content"] = current_text.strip()
        yield "", history
        time.sleep(0.08)

def start_training():
    return andrey.train(epochs=50)

with gr.Blocks(title="Андрей AI") as demo:
    gr.Markdown("# 🤖 Андрей AI (Оптимизированный)\n### Обучение на CPU (0 секунд лимита), генерация на ZeroGPU")
    
    chatbot = gr.Chatbot(height=400, label="Диалог")
    msg = gr.Textbox(label="Сообщение", placeholder="Напишите что-нибудь...")
    
    with gr.Row():
        clear = gr.Button("🧹 Очистить историю")
        train_btn = gr.Button("🎓 Обучить модель (50 эпох, ~3 сек)")
    
    status = gr.Textbox(label="Статус обучения", interactive=False)
    
    msg.submit(gradio_chat, [msg, chatbot], [msg, chatbot])
    clear.click(lambda: [], None, chatbot)
    train_btn.click(start_training, None, status)

if __name__ == "__main__":
    andrey = AndreyAI(MODEL_BIN, MODEL_GGUF)
    loaded = andrey.load_weights()
    
    if not loaded:
        print("⚠️ Модель не найдена. Нажмите кнопку 'Обучить модель' в интерфейсе.")
    
    demo.launch()