import json
from transformers import AutoTokenizer

MODEL_ID = "google/gemma-3-270m-it"
DATASET_PATH = "train.jsonl"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
token_lengths = []

print(f"Loading dataset from {DATASET_PATH}...")
with open(DATASET_PATH, 'r', encoding='utf-8') as f:
    for line in f:
        entry = json.loads(line)
        text = entry.get("text", "")
        # Tokenize the text and get the number of tokens
        length = len(tokenizer(text).input_ids)
        token_lengths.append(length)

print(f"Number of examples: {len(token_lengths)}")
if token_lengths:
    print(f"Min length:    {min(token_lengths)}")
    print(f"Max length:    {max(token_lengths)}")
    print(f"Average length:  {sum(token_lengths) / len(token_lengths):.2f}")