Update app.py

app.py

@@ -7,14 +7,17 @@ import re
 import gradio as gr
 from datetime import datetime
 import math
-import spaces  # <--- ВАЖНО: Импортируем библиотеку spaces
+import spaces  # Обязательно для ZeroGPU
 
-# ============ НАСТРОЙКИ ПУТЕЙ ============
+# ============ НАСТРОЙКИ ============
 DATA_DIR = '/data'
 os.makedirs(DATA_DIR, exist_ok=True)
-MODEL_PATH = os.path.join(DATA_DIR, 'pytorch_model_andrey_v6.bin')
+MODEL_PATH = os.path.join(DATA_DIR, 'andrey_zerogpu_v8.bin')
 
-DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+# На ZeroGPU мы не можем полагаться на глобальный DEVICE при инициализации,
+# поэтому будем определять его внутри функций, обернутых в @spaces.GPU
+# или использовать CPU для легких операций, если нужно.
+# Но для модели нужен CUDA.
 
 # ============ СЛОВАРЬ ============
 WORDS = [
@@ -71,8 +74,7 @@ def pad_sequence(seq, max_len=MAX_LEN):
     if len(seq) >= max_len: return seq[:max_len]
     return seq + [PAD] * (max_len - len(seq))
 
-# ============ TRANSFORMER АРХИТЕКТУРА ============
-
+# ============ МОДЕЛЬ ============
 class PositionalEncoding(nn.Module):
     def __init__(self, d_model, dropout=0.1, max_len=5000):
         super(PositionalEncoding, self).__init__()
@@ -108,19 +110,13 @@ class AndreyTransformer(nn.Module):
             dropout=dropout, batch_first=True
         )
         self.fc_out = nn.Linear(d_model, vocab_size)
-        self._init_weights()
-        
-    def _init_weights(self):
-        initrange = 0.1
-        self.embedding.weight.data.uniform_(-initrange, initrange)
-        self.fc_out.bias.data.zero_()
-        self.fc_out.weight.data.uniform_(-initrange, initrange)
         
-    def generate_mask(self, tgt_len):
-        return torch.triu(torch.ones(tgt_len, tgt_len), diagonal=1).bool().to(DEVICE)
+    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):
-        return (seq == pad_idx).to(DEVICE)
+    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):
@@ -135,7 +131,7 @@ class AndreyTransformer(nn.Module):
     
     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)
+        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):
@@ -146,7 +142,7 @@ class AndreyTransformer(nn.Module):
         )
         return self.fc_out(output)
 
-# ============ ДИАЛОГИ ============
+# ============ ДАННЫЕ ============
 DIALOGUES = [
     ("привет", "привет как дела"), ("здравствуй", "здравствуй рад тебя видеть"),
     ("доброе утро", "доброе утро хорошего дня"), ("добрый день", "добрый день чем помочь"),
@@ -271,59 +267,45 @@ class AndreyAI:
         self.bin_file = bin_file
         self.memory = {'chat_history': [], 'epochs_trained': 0}
         self.model = None
-        self.load()
-    
-    def _get_state(self):
-        return {
-            'model_state': self.model.state_dict() if self.model else None,
+        # Загружаем структуру, но веса загрузим позже
+        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_encoder_layers': 2, 'num_decoder_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': '7.0-ZeroGPU',
+            'memory': self.memory, 'version': '8.0-ZeroGPU-Fixed',
             'created': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         }
-    
-    def _restore(self, data):
-        self.memory = data.get('memory', self.memory)
-        if data.get('model_state'):
-            self.model = AndreyTransformer(
-                vocab_size=data.get('vocab_size', vocab_size),
-                d_model=data.get('d_model', 128), nhead=data.get('nhead', 4),
-                num_encoder_layers=data.get('num_encoder_layers', 2),
-                num_decoder_layers=data.get('num_decoder_layers', 2),
-                dim_feedforward=data.get('dim_feedforward', 256)
-            )
-            self.model.load_state_dict(data['model_state'])
-            self.model.to(DEVICE).eval()
-            return True
-        return False
-    
-    def load(self):
-        if os.path.exists(self.bin_file):
-            try:
-                if self._restore(torch.load(self.bin_file, map_location=DEVICE)):
-                    print(f"✅ Андрей загружен из {self.bin_file}")
-                    return True
-            except Exception as e: print(f"⚠️ Ошибка загрузки: {e}")
-        print("📝 Создаю нового Андрея...")
-        self.model = AndreyTransformer().to(DEVICE).eval()
-        return False
-    
-    def save(self):
-        os.makedirs(os.path.dirname(self.bin_file), exist_ok=True)
-        torch.save(self._get_state(), self.bin_file)
-        print(f"✅ Сохранён в /data: {os.path.getsize(self.bin_file)/1024:.1f} КБ")
-    
-    @spaces.GPU(duration=120) # <--- ВАЖНО: Декоратор для ZeroGPU
-    def train(self, epochs=150):
-        print("="*60)
-        print(f"🧠 ОБУЧЕНИЕ АНДРЕЯ — {epochs} ЭПОХ")
-        print("="*60)
+        torch.save(state, self.bin_file)
+        print(f"✅ Сохранено в /data: {os.path.getsize(self.bin_file)/1024:.1f} КБ")
+
+    @spaces.GPU(duration=120)
+    def train(self, epochs=50): # Уменьшил эпохи для быстроты теста, можно вернуть 150
+        print("🚀 Начало обучения на ZeroGPU...")
+        device = torch.device('cuda')
+        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)
-        self.model.train()
         
         for epoch in range(epochs):
             total_loss, n_batches = 0, 0
@@ -331,13 +313,13 @@ class AndreyAI:
             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)
+                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))
-                src_pad = self.model.create_pad_mask(src)
-                tgt_pad = self.model.create_pad_mask(tgt_in)
+                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)
@@ -349,12 +331,13 @@ class AndreyAI:
                 total_loss += loss.item()
                 n_batches += 1
             
-            if epoch % 10 == 0: print(f"Эпоха {epoch:3d}/{epochs} | Потери: {total_loss/n_batches:.4f}")
+            if epoch % 10 == 0: 
+                print(f"Эпоха {epoch}/{epochs} | Loss: {total_loss/n_batches:.4f}")
         
         self.memory['epochs_trained'] = epochs
-        print("\n✅ Обучение готово!")
-        self.save()
-        self.model.eval()
+        self.model.cpu() # Возвращаем на CPU для сохранения
+        self.save_weights()
+        print("✅ Обучение завершено!")
 
     def get_fallback_answer(self, question):
         q_clean = question.lower().strip()
@@ -363,7 +346,12 @@ class AndreyAI:
             if q in q_clean or q_clean in q: return a
         return "Интересный вопрос! Я еще учусь."
 
+    @spaces.GPU(duration=10) # ВАЖНО: Генерация тоже должна быть на GPU
     def generate(self, question, history=None, temperature=0.6, max_length=15):
+        device = torch.device('cuda')
+        self.model.to(device)
+        self.model.eval()
+        
         q = question.lower().strip()
         context_parts = []
         if history:
@@ -374,20 +362,24 @@ class AndreyAI:
         full_context = " ".join(context_parts)
         
         ctx_tokens = tokenize(full_context)
-        if not ctx_tokens: return self.get_fallback_answer(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)
+        # Создаем тензор уже на устройстве
+        src = torch.tensor([pad_sequence(ctx_tokens, MAX_LEN)], dtype=torch.long).to(device)
         generated_text = ""
         
         try:
             with torch.no_grad():
                 memory = self.model.encode(src)
                 response_tokens = []
-                decoder_input = torch.tensor([[START]], dtype=torch.long).to(DEVICE)
+                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)).to(DEVICE)
+                    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)
@@ -396,36 +388,37 @@ class AndreyAI:
                     
                     if next_token in [PAD, UNK]: break
                     response_tokens.append(next_token)
-                    decoder_input = torch.cat([decoder_input, torch.tensor([[next_token]], dtype=torch.long).to(DEVICE)], dim=1)
+                    decoder_input = torch.cat([decoder_input, torch.tensor([[next_token]], dtype=torch.long).to(device)], dim=1)
                 
                 generated_text = detokenize(response_tokens)
-        except Exception: pass
+        except Exception as e:
+            print(f"Ошибка генерации: {e}")
+        finally:
+            self.model.cpu() # Всегда возвращаем на CPU после работы
 
         return generated_text if generated_text else self.get_fallback_answer(q)
 
-# ============ GRADIO ИНТЕРФЕЙС ============
-def gradio_chat(question, history):
-    if not question: return "", history
-    answer = andrey.generate(question, history=history)
-    new_history = history + [(question, answer)]
-    andrey.memory['chat_history'] = new_history
+# ============ GRADIO ============
+def gradio_chat(message, history):
+    if not message: return "", history
+    answer = andrey.generate(message, history=history)
+    new_history = history + [(message, answer)]
     return "", new_history
 
-def launch_gradio():
-    # Исправлено для Gradio 6.0: theme передается в launch
-    with gr.Blocks(title="Андрей AI") as demo:
-        gr.Markdown("# 🤖 Андрей AI v7.0 (ZeroGPU)\n### Transformer с реальной памятью")
-        chatbot = gr.Chatbot(height=400, label="Диалог")
-        msg = gr.Textbox(label="Сообщение", placeholder="Напишите что-нибудь...")
-        clear = gr.Button("🧹 Очистить")
-        
-        msg.submit(gradio_chat, [msg, chatbot], [msg, chatbot])
-        clear.click(lambda: [], None, chatbot)
+with gr.Blocks(title="Андрей AI") as demo:
+    gr.Markdown("# 🤖 Андрей AI (ZeroGPU)\n### Transformer с памятью")
+    chatbot = gr.Chatbot(height=400, label="Диалог")
+    msg = gr.Textbox(label="Сообщение", placeholder="Напишите что-нибудь...")
+    clear = gr.Button("🧹 Очистить")
     
-    demo.launch(theme=gr.themes.Soft()) # <--- Исправлено место передачи темы
+    msg.submit(gradio_chat, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: [], None, chatbot)
 
 if __name__ == "__main__":
     andrey = AndreyAI(MODEL_PATH)
-    if andrey.model is None or andrey.memory.get('epochs_trained', 0) == 0:
-        andrey.train(150)
-    launch_gradio()
+    loaded = andrey.load_weights()
+    
+    if not loaded or andrey.memory.get('epochs_trained', 0) == 0:
+        andrey.train(150) # Обучаем если нет весов
+    
+    demo.launch()