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()