Update app.py

app.py

@@ -1,20 +1,23 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import gradio as gr
-import os
+import time
+import os # Добавлено для проверки наличия файла
 
 # --- КОНФИГУРАЦИЯ ---
+FILE_NAME = 'book.txt'
+MODEL_PATH = 'minigpt_checkpoint.pt'
 BLOCK_SIZE = 64
+BATCH_SIZE = 16
 EMBED_SIZE = 64
 HEADS = 4
-MODEL_PATH = 'minigpt_checkpoint.pt'
-FILE_NAME = 'book.txt'
+LR = 0.001
+EPOCHS = 300 # Увеличено для лучшего обучения на новых данных
 
-# --- АРХИТЕКТУРА ---
+# --- 1. АРХИТЕКТУРА МОДЕЛИ ---
 class MiniGPT(nn.Module):
     def __init__(self, vocab_size, embed_size, num_heads, block_size):
-        super().__init__()
+        super(MiniGPT, self).__init__()
         self.block_size = block_size
         self.embedding = nn.Embedding(vocab_size, embed_size)
         self.pos_embedding = nn.Embedding(block_size, embed_size)
@@ -25,75 +28,61 @@ class MiniGPT(nn.Module):
     def forward(self, x):
         B, T = x.shape
         pos = torch.arange(T, device=x.device).unsqueeze(0)
-        out = self.embedding(x) + self.pos_embedding(pos)
+        tok_emb = self.embedding(x)
+        pos_emb = self.pos_embedding(pos)
+        out = tok_emb + pos_emb
         out = self.transformer(out)
-        return self.fc_out(out)
+        logits = self.fc_out(out)
+        return logits
 
-# --- ДАННЫЕ И ТОКЕНИЗАЦИЯ ---
-if os.path.exists(FILE_NAME):
-    with open(FILE_NAME, 'r', encoding='utf-8') as f: text = f.read()
-else:
-    text = "привет как дела нормально пока" * 100
+# --- 2. ПОДГОТОВКА ДАННЫХ И ТОКЕНИЗАЦИЯ ---
+try:
+    with open(FILE_NAME, 'r', encoding='utf-8') as f:
+        text = f.read()
+    print(f"Успешно прочитан файл: {FILE_NAME}, размер текста: {len(text)} символов.")
+except FileNotFoundError:
+    print(f"Ошибка: файл '{FILE_NAME}' не найден. Использую fallback текст.")
+    # Fallback текст должен содержать символы '<', '|', '>', 'u', 's', 'e', 'r', 'm', 'o', 'd', 'l'
+    text = "<|user|>привет<|model|>нормально" * 100
 
 chars = sorted(list(set(text)))
 vocab_size = len(chars)
 stoi = { ch:i for i,ch in enumerate(chars) }
 itos = { i:ch for i,ch in enumerate(chars) }
-encode = lambda s: [stoi.get(c, 0) for c in s]
+encode = lambda s: [stoi[c] for c in s]
 decode = lambda l: ''.join([itos[i] for i in l])
 
-# --- ЗАГРУЗКА МОДЕЛИ ---
+data = torch.tensor(encode(text), dtype=torch.long)
+print(f"Данные закодированы в тензор размером: {data.shape}")
+
+# --- 3. НАСТРОЙКИ ОБУЧЕНИЯ И ИНИЦИАЛИЗАЦИЯ ---
 model = MiniGPT(vocab_size, EMBED_SIZE, HEADS, BLOCK_SIZE)
-if os.path.exists(MODEL_PATH):
-    # Загружаем модель на CPU, что важно для HF Spaces с базовым тарифом
-    model.load_state_dict(torch.load(MODEL_PATH, map_location=torch.device('cpu')))
-model.eval()
+optimizer = torch.optim.Adam(model.parameters(), lr=LR)
+criterion = nn.CrossEntropyLoss()
+
+# --- 4. ЦИКЛ ОБУЧЕНИЯ ---
+print("Начинаю обучение...")
+model.train()
+
+for epoch in range(EPOCHS):
+    # Генерация случайных батчей из данных
+    ix = torch.randint(len(data) - BLOCK_SIZE, (BATCH_SIZE,))
+    xb = torch.stack([data[i:i+BLOCK_SIZE] for i in ix])
+    yb = torch.stack([data[i+1:i+BLOCK_SIZE+1] for i in ix])
 
-# --- ЛОГИКА ГЕНЕРАЦИИ С ТОКЕНАМИ ---
-def predict(prompt, max_length, temperature):
-    if not prompt: return "Введите текст"
-    
-    # Принудительно добавляем токен Модели к запросу пользователя
-    full_prompt = prompt.strip() + "<|model|>"
-    
-    context_tokens = encode(full_prompt)[-BLOCK_SIZE:]
-    context = torch.tensor(context_tokens, dtype=torch.long).unsqueeze(0)
-    
-    generated_tokens = []
-    for _ in range(max_length):
-        cond = context[:, -BLOCK_SIZE:]
-        with torch.no_grad():
-            logits = model(cond)[:, -1, :]
-            
-            if temperature == 0:
-                probs = F.softmax(logits, dim=-1)
-                next_token = torch.argmax(probs, dim=-1).unsqueeze(0)
-            else:
-                probs = F.softmax(logits / temperature, dim=-1)
-                next_token = torch.multinomial(probs, num_samples=1)
-            
-            # Остановка генерации, если модель сгенерировала начало токена '<'
-            if decode([next_token.item()]) == '<':
-                 break
+    logits = model(xb)
+    B, T, C = logits.shape
+    loss = criterion(logits.view(B*T, C), yb.view(B*T))
 
-            context = torch.cat((context, next_token), dim=1)
-            generated_tokens.append(next_token.item())
-            
-    return decode(generated_tokens)
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer.step()
 
-# --- ИНТЕРФЕЙС GRADIO ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🤖 MiniGPT Chat с настройками")
-    with gr.Row():
-        with gr.Column():
-            input_text = gr.Textbox(label="Ваш запрос (начинайте с <|user|>)", placeholder="Напишите начало фразы...", lines=3)
-            max_len_slider = gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Максимальная длина ответа")
-            temp_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="Температура (0=детерминированный)")
-            btn = gr.Button("Сгенерировать")
+    if epoch % 50 == 0:
+        print(f"Эпоха {epoch}, Ошибка: {loss.item():.4f}")
 
-        output_text = gr.Textbox(label="Ответ модели", lines=10)
-    
-    btn.click(fn=predict, inputs=[input_text, max_len_slider, temp_slider], outputs=[output_text])
+print("Обучение завершено.")
 
-if __name__ == "__main__":
-    demo.launch()
+# --- 5. СОХРАНЕНИЕ МОДЕЛИ ---
+torch.save(model.state_dict(), MODEL_PATH)
+print(f"Модель сохранена в файл {MODEL_PATH}")