Update README.md

README.md

@@ -1,51 +1,115 @@
 ```python
-if os.path.exists(checkpoint_path):
-    checkpoint = torch.load(checkpoint_path)
+embedding_layer = embedding_layer.to(DEVICE)
+transformer_encoder = transformer_encoder.to(DEVICE)
+pos_encoding = pos_encoding.to(DEVICE)
+output_layer = output_layer.to(DEVICE)
+
+# -----------------------------
+# Оптимизатор
+# -----------------------------
+optimizer = torch.optim.Adam(
+    list(embedding_layer.parameters()) +
+    list(transformer_encoder.parameters()) +
+    list(pos_encoding.parameters()) +
+    list(output_layer.parameters()),
+    lr=1e-4
+)
+
+# -----------------------------
+# Загружаем чекпоинт
+# -----------------------------
+start_epoch = 0
+if os.path.exists(CHECKPOINT_PATH):
+    checkpoint = torch.load(CHECKPOINT_PATH, map_location=DEVICE)
     embedding_layer.load_state_dict(checkpoint['embedding_state'])
-    pos_encoding.load_state_dict(checkpoint['pos_encoding_state']) 
-    transformer_encoderLayer.load_state_dict(checkpoint['transformer_state'])
+    pos_encoding.load_state_dict(checkpoint['pos_encoding_state'])
+    transformer_encoder.load_state_dict(checkpoint['transformer_state'])
     output_layer.load_state_dict(checkpoint['output_state'])
     optimizer.load_state_dict(checkpoint['optimizer_state'])
     start_epoch = checkpoint['epoch'] + 1
-    print(f" Модель загружена, продолжаем с эпохи {start_epoch}")
+    print(f"Модель загружена, продолжаем с эпохи {start_epoch}")
 else:
-    start_epoch = 0
-    print(" Чекпоинт не найден, начинаем обучение с нуля")
-
-epochNum = 10
-for epoch in range(epochNum):
-    optimizer.zero_grad()
-    epochmy = start_epoch + epoch
-    embedded = embedding_layer(input_ids)
-    embedded = pos_encoding(embedded)
-    src = embedded.transpose(0, 1)
-
-    outputTransformer = transformer_encoderLayer(src, src_key_padding_mask=(attention_mask == 0))
-    outputTransformer = outputTransformer.transpose(0, 1)  # обратно [batch, seq_len, embedding_dim]
-
-    logits = output_layer(outputTransformer)
-    loss = criterion(logits.view(-1, vocab_size), target_ids.view(-1))
-    before = pos_encoding.pos_embedding.weight.clone()
-    loss.backward()
-    optimizer.step()  # обновляем веса
-    after = pos_encoding.pos_embedding.weight
-    print(f"Изменение весов pos_encoding: {(after - before).abs().sum():.6f}")
-    print("Loss:", loss.item())
-
-    # После обучения (или внутри цикла, чтобы смотреть динамику)
-    with torch.no_grad():
-        embedded = embedding_layer(input_ids)
-        embedded = pos_encoding(embedded)
-        src = embedded.transpose(0, 1)
-        outputTransformer = transformer_encoderLayer(src, src_key_padding_mask=(attention_mask == 0))
-        outputTransformer = outputTransformer.transpose(0, 1)
-        logits = output_layer(outputTransformer)  # [batch, seq_len, vocab_size]
-
-        # Берём самый вероятный токен для каждого положения
-        predicted_token_ids = torch.argmax(logits, dim=-1)  # [batch, seq_len]
-
-        # Переводим индексы обратно в текст
-        predicted_text = tokenizer.batch_decode(predicted_token_ids, skip_special_tokens=False)
-        print("Predicted text:", predicted_text[0])
-
-        print("Loss before backward:", loss.item())
+    print("Чекпоинт не найден, начинаем обучение с нуля")
+
+# -----------------------------
+# Обучение с отладкой
+# -----------------------------
+for epoch in range(start_epoch, start_epoch + EPOCHS):
+    running_loss = 0.0
+    print(f"\n=== Эпоха {epoch + 1}/{start_epoch + EPOCHS} ===")
+
+    for chunk_idx, (input_ids_chunk, attention_mask_chunk, target_ids_chunk) in enumerate(
+            chunked_tokenizer(data, tokenizer, max_len=MAX_LEN, chunk_size=CHUNK_SIZE)
+    ):
+        print(f"\n--- Чанк {chunk_idx + 1} / {len(data) // CHUNK_SIZE + 1} ---")
+        dataset = TensorDataset(input_ids_chunk, attention_mask_chunk, target_ids_chunk)
+        dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+
+        for batch_idx, batch in enumerate(dataloader):
+            batch_input_ids, batch_attention_mask, batch_target_ids = [x.to(DEVICE) for x in batch]
+            padding_mask = (batch_attention_mask == 0)
+
+            optimizer.zero_grad()
+
+            # Эмбеддинги
+            embedded = embedding_layer(batch_input_ids)
+            print(f"[DEBUG] embedded shape: {embedded.shape}")  # batch, seq_len, embed_dim
+
+            # Позиционное кодирование
+            embedded = embedded.transpose(0, 1)  # seq_len, batch, embed_dim
+            embedded = pos_encoding(embedded)
+            print(f"[DEBUG] embedded + pos_encoding shape: {embedded.shape}")
+
+            # Трансформер
+            transformer_output = transformer_encoder(embedded, src_key_padding_mask=padding_mask)
+            transformer_output = transformer_output.transpose(0, 1)  # batch, seq_len, embed_dim
+            print(f"[DEBUG] transformer_output shape: {transformer_output.shape}")
+
+            # Память выхода трансформера (примерно)
+            batch_size, seq_len, emb_dim = transformer_output.shape
+            mem_MB = batch_size * seq_len * emb_dim * 4 / 1024 ** 2
+            print(f"[DEBUG] Output memory approx: {mem_MB:.2f} MB")
+
+            # Линейный слой
+            logits = output_layer(transformer_output)
+            print(f"[DEBUG] logits shape: {logits.shape}")
+
+            # Потери
+            loss = criterion(logits.view(-1, vocab_size), batch_target_ids.view(-1))
+            loss_history.append(loss.item())
+            print(f"[DEBUG] batch {batch_idx + 1} loss: {loss.item():.6f}")
+
+            # Backprop
+            loss.backward()
+            optimizer.step()
+
+            running_loss += loss.item() * batch_input_ids.size(0)
+
+            # Демонстрация предсказаний
+            pred_tokens = torch.argmax(logits, dim=-1)
+            sample_input = tokenizer.decode(batch_input_ids[0], skip_special_tokens=True)
+            sample_pred = tokenizer.decode(pred_tokens[0], skip_special_tokens=True)
+            sample_target = tokenizer.decode(batch_target_ids[0], skip_special_tokens=True)
+            print(f"[DEBUG] Sample input:  {sample_input[:50]}...")
+            print(f"[DEBUG] Sample target: {sample_target[:50]}...")
+            print(f"[DEBUG] Sample pred:   {sample_pred[:50]}...")
+
+            # Очистка памяти
+            del batch_input_ids, batch_attention_mask, batch_target_ids, embedded, transformer_output, logits
+            torch.cuda.empty_cache()
+
+    avg_loss = running_loss / len(data)
+    print(f"\n=== Эпоха {epoch + 1} завершена — Avg Loss: {avg_loss:.6f} ===\n")
+
+# -----------------------------
+# Сохраняем чекпоинт
+# -----------------------------
+torch.save({
+    'embedding_state': embedding_layer.state_dict(),
+    'pos_encoding_state': pos_encoding.state_dict(),
+    'transformer_state': transformer_encoder.state_dict(),
+    'output_state': output_layer.state_dict(),
+    'optimizer_state': optimizer.state_dict(),
+    'epoch': epoch
+}, CHECKPOINT_PATH)
+