import os
from pathlib import Path
import torch
from datasets import load_dataset
from transformers import T5Tokenizer, T5ForConditionalGeneration
from datetime import datetime
from torch.utils.data import IterableDataset, DataLoader
from tqdm import tqdm

class CustomIterableDataset(IterableDataset):
    def __init__(self, dataset):
        self.dataset = dataset

    def __iter__(self):
        for example in self.dataset:
            input_ids = torch.tensor(example["input_ids"], dtype=torch.long)
            attention_mask = torch.tensor(example.get("attention_mask", [1] * len(example["input_ids"])), dtype=torch.long)
            labels = torch.tensor(example["labels"], dtype=torch.long)

            yield {
                "input_ids": input_ids,
                "attention_mask": attention_mask,
                "labels": labels
            }

def main():
    base_path = Path("E:/aiv2")
    cache_path = base_path / 'models'

    os.environ['HF_HOME'] = str(cache_path / 'huggingface')
    os.environ['HF_DATASETS_CACHE'] = str(cache_path / 'datasets_cache')
    os.environ['PYARROW_TMPDIR'] = str(cache_path / 'temp')
    os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

    # Twórz katalogi, jeśli nie istnieją
    cache_dirs = [cache_path / 'huggingface', cache_path / 'datasets_cache', cache_path / 'temp']
    for dir_path in cache_dirs:
        dir_path.mkdir(parents=True, exist_ok=True)

    # Ustawienie urządzenia na CPU dla kroków debugowania
    device = torch.device('cpu')
    print(f"Using CPU for computation.")

    # Inicjalizacja tokenizera
    tokenizer = T5Tokenizer.from_pretrained('t5-small', cache_dir=str(cache_path / "huggingface"), use_fast=True)

    # Ładowanie i przetwarzanie zbioru danych
    try:
        dataset = load_dataset("neuralwork/arxiver", split='train', streaming=True)
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return

    # Identyfikacja używanej kolumny tekstowej
    text_column = None
    try:
        sample_entry = next(iter(dataset))
    except Exception as e:
        print(f"Error iterating over dataset: {e}")
        return

    for possible_column in ['content', 'text', 'title']:
        if possible_column in sample_entry:
            text_column = possible_column
            break

    if text_column is None:
        raise KeyError("Couldn't find a suitable text column.")

    print(f"Using text column: {text_column}")

    # Funkcja przygotowująca dane
    def preprocess_function(examples):
        inputs = examples[text_column]
        model_inputs = tokenizer(inputs, max_length=256, truncation=True, padding="max_length")
        labels = model_inputs["input_ids"].copy()
        model_inputs["labels"] = [
            [label if label != tokenizer.pad_token_id else -100 for label in label_ids]
            for label_ids in labels
        ]
        return model_inputs

    # Przetwarzanie zbioru danych
    try:
        train_dataset = dataset.map(preprocess_function, batched=True, batch_size=100, remove_columns=dataset.column_names)
    except Exception as e:
        print(f"Error processing dataset: {e}")
        return

    # Konfiguracja DataLoader
    try:
        train_pytorch_dataset = CustomIterableDataset(train_dataset)
        train_dataloader = DataLoader(train_pytorch_dataset, batch_size=8, num_workers=0)
    except Exception as e:
        print(f"Error creating DataLoader: {e}")
        return

    # Ładowanie modelu
    try:
        model = T5ForConditionalGeneration.from_pretrained('t5-small', cache_dir=str(cache_path / "huggingface")).to(device)
        print("Model loaded successfully on CPU.")
    except Exception as e:
        print(f"Error loading model: {e}")
        return

    # Konfiguracja optymalizatora
    optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5)

    # Trenowanie modelu
    model.train()
    num_epochs = 3

    for epoch in range(num_epochs):
        print(f"Starting epoch {epoch + 1}/{num_epochs}")
        epoch_loss = 0
        batch_count = 0

        for batch in tqdm(train_dataloader, desc=f"Epoch {epoch + 1}"):
            optimizer.zero_grad()

            try:
                input_ids = batch["input_ids"].to(device)
                attention_mask = batch["attention_mask"].to(device)
                labels = batch["labels"].to(device)

                outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                loss = outputs.loss
                loss.backward()
                optimizer.step()
            except Exception as e:
                print(f"Error during forward/backward pass: {e}")
                print("Inspecting tensor values:")
                print(f"input_ids sample: {input_ids[0][:10]}")
                print(f"attention_mask sample: {attention_mask[0][:10]}")
                print(f"labels sample: {labels[0][:10]}")
                continue

            epoch_loss += loss.item()
            batch_count += 1

            if batch_count % 10 == 0:
                print(f"Epoch: {epoch + 1}, Batch: {batch_count}, Loss: {loss.item():.4f}")

        avg_loss = epoch_loss / batch_count if batch_count > 0 else 0
        print(f"Epoch {epoch + 1} ended. Average Loss: {avg_loss:.4f}")

    # Zapis modelu i tokenizera
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    save_dir = cache_path / f'T5_LANGAL_{timestamp}'
    save_dir.mkdir(parents=True, exist_ok=True)

    try:
        model.save_pretrained(str(save_dir))
        tokenizer.save_pretrained(str(save_dir))
        print(f"Model and tokenizer saved in {save_dir}")
    except Exception as e:
        print(f"Error saving model/tokenizer: {e}")

if __name__ == "__main__":
    main() napraw i doprowadź do dziłania kodu możesz go przetestowac