from fastapi import FastAPI
from contextlib import asynccontextmanager
from transformers import PreTrainedTokenizerFast
from tokenizers import ByteLevelBPETokenizer
from datasets import load_dataset

app = FastAPI()

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Εκκίνηση διαδικασίας εκπαίδευσης tokenizer
    await train_tokenizer()
    yield
    # Οποιαδήποτε διαδικασία κατά το κλείσιμο της εφαρμογής

app = FastAPI(lifespan=lifespan)

async def train_tokenizer():
    # Μέγεθος λεξιλογίου και ελάχιστη συχνότητα
    vocab_size = 50000
    min_frequency = 2

    # Φόρτωση δεδομένων από Oscar και Wikipedia μέσω streaming
    dataset_greek = load_dataset("oscar", "unshuffled_deduplicated_el", split="train", streaming=True)
    dataset_english = load_dataset("wikipedia", "20220301.en", split="train", streaming=True)

    # Διαχείριση καθαρού κώδικα (π.χ., από GitHub ή άλλες πηγές αν υπάρχουν)
    try:
        dataset_code = load_dataset("bigcode/the-stack", split="train", streaming=True)
        datasets_list = [dataset_greek, dataset_english, dataset_code]
    except:
        datasets_list = [dataset_greek, dataset_english]

    # Ενοποίηση των δεδομένων και προεπεξεργασία
    def preprocess_data(dataset):
        for item in dataset:
            text = item["text"]
            # Αφαίρεση περιττών χαρακτήρων και κανονικοποίηση
            text = text.strip().lower()  # Μπορείς να το κάνεις lower εάν το θέλεις
            if text:  # Εξασφαλίζουμε ότι δεν είναι άδειο το κείμενο
                yield text

    combined_data = (
        preprocess_data(dataset) for dataset in datasets_list
    )

    # Δημιουργία του tokenizer
    tokenizer = ByteLevelBPETokenizer()

    # Εκπαίδευση του tokenizer
    tokenizer.train_from_iterator(
        combined_data,
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
    )

    # Αποθήκευση του tokenizer
    tokenizer.save_model(".")

@app.get("/")
async def root():
    return {"message": "Custom Tokenizer Training Completed and Saved"}