Update app.py

app.py

@@ -7,7 +7,7 @@ import re
 import gradio as gr
 from datetime import datetime
 import math
-import numpy as np
+import spaces  # <--- ВАЖНО: Импортируем библиотеку spaces
 
 # ============ НАСТРОЙКИ ПУТЕЙ ============
 DATA_DIR = '/data'
@@ -39,7 +39,6 @@ WORDS = [
     'работа', 'отдых', 'путешествие', 'еда', 'вода', 'спорт', 'люди', 'мир', 'знание',
     'будущее', 'прошлое', 'настоящее', 'интерес', 'радость', 'успех', 'дружба', 'любовь',
     'семья', 'здоровье', 'счастье', 'удача', 'смех', 'солнце', 'звезды', 'мечта',
-    # Добавляем маркеры ролей для памяти
     'пользователь', 'андрей', 'говорит' 
 ]
 
@@ -53,7 +52,7 @@ vocab_size = len(WORDS) + 3
 PAD = 0
 UNK = 1
 START = 2
-MAX_LEN = 30  # Увеличиваем длину, чтобы влезала история
+MAX_LEN = 30
 
 def tokenize(text):
     return [word_to_idx.get(w, UNK) for w in text.lower().split()]
@@ -61,20 +60,15 @@ def tokenize(text):
 def detokenize(tokens):
     words = []
     for t in tokens:
-        if t == START:
-            continue
-        if t == PAD:
-            break
-        if t == UNK:
-            continue 
+        if t == START: continue
+        if t == PAD: break
+        if t == UNK: continue 
         w = idx_to_word.get(t)
-        if w:
-            words.append(w)
+        if w: words.append(w)
     return ' '.join(words)
 
 def pad_sequence(seq, max_len=MAX_LEN):
-    if len(seq) >= max_len:
-        return seq[:max_len]
+    if len(seq) >= max_len: return seq[:max_len]
     return seq + [PAD] * (max_len - len(seq))
 
 # ============ TRANSFORMER АРХИТЕКТУРА ============
@@ -83,11 +77,9 @@ class PositionalEncoding(nn.Module):
     def __init__(self, d_model, dropout=0.1, max_len=5000):
         super(PositionalEncoding, self).__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)
         pe = pe.unsqueeze(0)
@@ -102,24 +94,19 @@ class AndreyTransformer(nn.Module):
                  num_encoder_layers=2, num_decoder_layers=2, 
                  dim_feedforward=256, dropout=0.1, max_len=MAX_LEN):
         super().__init__()
-        
         self.d_model = d_model
         self.vocab_size = vocab_size
-        
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.pos_encoder = PositionalEncoding(d_model, dropout, max_len)
         self.pos_decoder = PositionalEncoding(d_model, dropout, max_len)
         
         self.transformer = nn.Transformer(
-            d_model=d_model,
-            nhead=nhead,
+            d_model=d_model, nhead=nhead,
             num_encoder_layers=num_encoder_layers,
             num_decoder_layers=num_decoder_layers,
             dim_feedforward=dim_feedforward,
-            dropout=dropout,
-            batch_first=True
+            dropout=dropout, batch_first=True
         )
-        
         self.fc_out = nn.Linear(d_model, vocab_size)
         self._init_weights()
         
@@ -130,8 +117,7 @@ class AndreyTransformer(nn.Module):
         self.fc_out.weight.data.uniform_(-initrange, initrange)
         
     def generate_mask(self, tgt_len):
-        mask = torch.triu(torch.ones(tgt_len, tgt_len), diagonal=1).bool()
-        return mask.to(DEVICE)
+        return torch.triu(torch.ones(tgt_len, tgt_len), diagonal=1).bool().to(DEVICE)
     
     def create_pad_mask(self, seq, pad_idx=PAD):
         return (seq == pad_idx).to(DEVICE)
@@ -140,36 +126,27 @@ class AndreyTransformer(nn.Module):
                 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_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
         )
-        
-        output = self.fc_out(output)
-        return output
+        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)
-        memory = self.transformer.encoder(src_emb, src_key_padding_mask=src_key_padding_mask)
-        return memory
+        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_emb, memory, tgt_mask=tgt_mask,
             tgt_key_padding_mask=tgt_key_padding_mask
         )
-        
         return self.fc_out(output)
 
-# ============ ДИАЛОГИ ДЛЯ ОБУЧЕНИЯ ============
+# ============ ДИАЛОГИ ============
 DIALOGUES = [
     ("привет", "привет как дела"), ("здравствуй", "здравствуй рад тебя видеть"),
     ("доброе утро", "доброе утро хорошего дня"), ("добрый день", "добрый день чем помочь"),
@@ -275,60 +252,35 @@ DIALOGUES = [
 
 FALLBACK_DICT = {q.lower(): a for q, a in DIALOGUES}
 
-# ============ ПОДГОТОВКА ДАННЫХ ============
 def prepare_data():
-    X_questions = []
-    Y_answers_input = []
-    Y_answers_target = []
-    
+    X_questions, Y_answers_input, Y_answers_target = [], [], []
     for q, a in DIALOGUES:
-        q_toks = tokenize(q)
-        a_toks = tokenize(a)
-        
+        q_toks, a_toks = tokenize(q), tokenize(a)
         if q_toks and a_toks:
-            q_padded = pad_sequence(q_toks, MAX_LEN)
-            
-            a_input = [START] + a_toks
-            a_input_padded = pad_sequence(a_input, MAX_LEN)
-            
-            a_target = a_toks + [PAD]
-            a_target_padded = pad_sequence(a_target, MAX_LEN)
-            
-            X_questions.append(q_padded)
-            Y_answers_input.append(a_input_padded)
-            Y_answers_target.append(a_target_padded)
-    
-    X_tensor = torch.tensor(X_questions, dtype=torch.long)
-    Y_input_tensor = torch.tensor(Y_answers_input, dtype=torch.long)
-    Y_target_tensor = torch.tensor(Y_answers_target, dtype=torch.long)
-    
-    print(f"📚 Всего примеров: {len(X_tensor)}")
-    return X_tensor, Y_input_tensor, Y_target_tensor
+            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))
 
-# ============ АНДРЕЙ TRANSFORMER ============
 class AndreyAI:
     def __init__(self, bin_file=MODEL_PATH):
         self.bin_file = bin_file
-        self.memory = {
-            'chat_history': [], # Здесь хранится реальная история
-            'epochs_trained': 0
-        }
+        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,
-            'vocab_size': vocab_size,
-            'd_model': 128,
-            'nhead': 4,
-            'num_encoder_layers': 2,
-            'num_decoder_layers': 2,
-            'dim_feedforward': 256,
+            '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': '6.0-Memory-Transformer',
+            'memory': self.memory, 'version': '7.0-ZeroGPU',
             'created': datetime.now().strftime("%Y-%m-%d %H:%M:%S")
         }
     
@@ -337,84 +289,59 @@ class AndreyAI:
         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),
+                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)
-            self.model.eval()
+            self.model.to(DEVICE).eval()
             return True
         return False
     
     def load(self):
         if os.path.exists(self.bin_file):
             try:
-                data = torch.load(self.bin_file, map_location=DEVICE)
-                if self._restore(data):
+                if self._restore(torch.load(self.bin_file, map_location=DEVICE)):
                     print(f"✅ Андрей загружен из {self.bin_file}")
-                    print(f"🧠 Обучен: {self.memory.get('epochs_trained', 0)} эпох")
-                    print(f"💾 История чатов: {len(self.memory.get('chat_history', []))} сообщений")
                     return True
-            except Exception as e:
-                print(f"⚠️ Ошибка загрузки: {e}")
-        
+            except Exception as e: print(f"⚠️ Ошибка загрузки: {e}")
         print("📝 Создаю нового Андрея...")
-        self.model = AndreyTransformer().to(DEVICE)
-        self.model.eval()
+        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)
-        size = os.path.getsize(self.bin_file) / 1024
-        print(f"✅ Сохранён в /data: {size:.1f} КБ")
+        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(f"📂 Модель: {self.bin_file}")
-        print("="*60 + "\n")
-        
+        print("="*60)
         X_data, Y_input, Y_target = prepare_data()
-        
         criterion = nn.CrossEntropyLoss(ignore_index=PAD)
-        optimizer = optim.Adam(self.model.parameters(), lr=0.0005)
-        scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
-        
+        optimizer = optim.Adam(self.model.parameters(), lr=0.001)
         self.model.train()
         
         for epoch in range(epochs):
-            total_loss = 0
-            n_batches = 0
-            
+            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_input = Y_input[idx].unsqueeze(0).to(DEVICE)
-                tgt_target = Y_target[idx].unsqueeze(0).to(DEVICE)
+                tgt_in = Y_input[idx].unsqueeze(0).to(DEVICE)
+                tgt_tar = Y_target[idx].unsqueeze(0).to(DEVICE)
                 
-                tgt_len = tgt_input.size(1)
-                tgt_mask = self.model.generate_mask(tgt_len)
-                
-                src_key_padding_mask = self.model.create_pad_mask(src)
-                tgt_key_padding_mask = self.model.create_pad_mask(tgt_input)
+                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)
                 
                 optimizer.zero_grad()
-                
-                output = self.model(
-                    src, tgt_input, 
-                    tgt_mask=tgt_mask,
-                    src_key_padding_mask=src_key_padding_mask,
-                    tgt_key_padding_mask=tgt_key_padding_mask
-                )
-                
-                loss = criterion(output.view(-1, vocab_size), tgt_target.view(-1))
-                
+                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()
@@ -422,175 +349,83 @@ class AndreyAI:
                 total_loss += loss.item()
                 n_batches += 1
             
-            scheduler.step()
-            
-            if epoch % 10 == 0:
-                avg_loss = total_loss / n_batches
-                print(f"Эпоха {epoch:3d}/{epochs} | Потери: {avg_loss:.4f}")
+            if epoch % 10 == 0: print(f"Эпоха {epoch:3d}/{epochs} | Потери: {total_loss/n_batches:.4f}")
         
         self.memory['epochs_trained'] = epochs
         print("\n✅ Обучение готово!")
         self.save()
         self.model.eval()
-    
+
     def get_fallback_answer(self, question):
         q_clean = question.lower().strip()
-        if q_clean in FALLBACK_DICT:
-            return FALLBACK_DICT[q_clean]
+        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
+            if q in q_clean or q_clean in q: return a
         return "Интересный вопрос! Я еще учусь."
 
     def generate(self, question, history=None, temperature=0.6, max_length=15):
-        """
-        Генерирует ответ с учетом истории (Real Memory).
-        history: список кортежей [(user_msg, bot_msg), ...]
-        """
         q = question.lower().strip()
-        
-        # 1. Формируем контекст из истории
         context_parts = []
         if history:
-            # Берем последние 3 пары сообщений, чтобы не превысить MAX_LEN
-            recent_history = history[-3:]
-            for user_msg, bot_msg in recent_history:
+            for user_msg, bot_msg in history[-3:]:
                 context_parts.append(f"пользователь говорит {user_msg}")
                 context_parts.append(f"андрей говорит {bot_msg}")
-        
-        # Добавляем текущий вопрос
         context_parts.append(f"пользователь говорит {question}")
-        
-        # Склеиваем в одну строку
         full_context = " ".join(context_parts)
         
-        # 2. Токенизация контекста
         ctx_tokens = tokenize(full_context)
-        if not ctx_tokens:
-            return self.get_fallback_answer(q)
-        
-        # Обрезаем до MAX_LEN, оставляя конец (самое важное - текущий вопрос)
-        if len(ctx_tokens) > MAX_LEN:
-            ctx_tokens = ctx_tokens[-MAX_LEN:]
-            
-        ctx_padded = pad_sequence(ctx_tokens, MAX_LEN)
-        src = torch.tensor([ctx_padded], dtype=torch.long).to(DEVICE)
+        if not ctx_tokens: 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)
         generated_text = ""
         
         try:
             with torch.no_grad():
                 memory = self.model.encode(src)
-                
                 response_tokens = []
                 decoder_input = torch.tensor([[START]], dtype=torch.long).to(DEVICE)
                 
                 for i in range(max_length):
-                    tgt_len = decoder_input.size(1)
-                    tgt_mask = self.model.generate_mask(tgt_len).to(DEVICE)
-                    
+                    tgt_mask = self.model.generate_mask(decoder_input.size(1)).to(DEVICE)
                     output = self.model.decode_step(decoder_input, memory, tgt_mask=tgt_mask)
-                    
                     logits = output[:, -1, :] / temperature
                     probs = torch.softmax(logits, dim=-1)
-                    
-                    top_prob, next_token = torch.max(probs, dim=-1)
+                    _, next_token = torch.max(probs, dim=-1)
                     next_token = next_token.item()
-                    confidence = top_prob.item()
-                    
-                    if next_token == PAD or next_token == UNK or confidence < 0.05:
-                        break
                     
+                    if next_token in [PAD, UNK]: break
                     response_tokens.append(next_token)
-                    next_token_tensor = torch.tensor([[next_token]], dtype=torch.long).to(DEVICE)
-                    decoder_input = torch.cat([decoder_input, next_token_tensor], 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 as e:
-            print(f"Ошибка генерации: {e}")
-            generated_text = ""
+        except Exception: pass
 
-        # 3. Fallback если модель промолчала
-        if not generated_text or len(generated_text.split()) < 1:
-            return self.get_fallback_answer(q)
-            
-        return generated_text
-    
-    def chat(self):
-        print("="*60)
-        print("🤖 АНДРЕЙ v6.0 (Transformer + Real Memory)")
-        print("💾 Память сохраняется в /data")
-        print("="*60 + "\n")
-        
-        history = self.memory.get('chat_history', [])
-        
-        while True:
-            user = input("👤 Вы: ").strip()
-            
-            if user.lower() in ['пока', 'выход', 'exit']:
-                print("🤖 Андрей: Пока! 👋")
-                self.save()
-                break
-            
-            if not user:
-                continue
-            
-            answer = self.generate(user, history=history)
-            print(f"🤖 Андрей: {answer}\n")
-            
-            # Обновляем историю
-            history.append((user, answer))
-            self.memory['chat_history'] = history
+        return generated_text if generated_text else self.get_fallback_answer(q)
 
 # ============ GRADIO ИНТЕРФЕЙС ============
 def gradio_chat(question, history):
-    if not question:
-        return "", history
-    
-    # Передаем текущую историю в модель
+    if not question: return "", history
     answer = andrey.generate(question, history=history)
-    
-    # Обновляем историю
     new_history = history + [(question, answer)]
-    
-    # Сохраняем обновленную историю в память объекта
     andrey.memory['chat_history'] = new_history
-    
     return "", new_history
 
 def launch_gradio():
-    with gr.Blocks(title="Андрей AI", theme=gr.themes.Soft()) as demo:
-        gr.Markdown("""
-        # 🤖 Андрей AI v6.0
-        ### Transformer с реальной памятью
-        **Память:** Хранится в `/data`
-        **Контекст:** Помнит последние сообщения
-        """)
-        
-        chatbot = gr.Chatbot(height=400, label="Диалог с Андреем")
-        msg = gr.Textbox(label="Ваше сообщение", placeholder="Напишите что-нибудь...")
-        clear = gr.Button("🧹 Очистить историю")
+    # Исправлено для 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)
-        
-        gr.Markdown("""
-        ### ❓ Примеры:
-        - Привет
-        - Меня зовут Евгений
-        - Как меня зовут? (должен вспомнить)
-        - 2+2
-        """)
     
-    demo.launch(share=True)
+    demo.launch(theme=gr.themes.Soft()) # <--- Исправлено место передачи темы
 
-# ============ ЗАПУСК ============
 if __name__ == "__main__":
     andrey = AndreyAI(MODEL_PATH)
-    
     if andrey.model is None or andrey.memory.get('epochs_trained', 0) == 0:
         andrey.train(150)
-    
-    print("\n🚀 Запуск Gradio интерфейса...")
     launch_gradio()