Function prepare_glove_dataset

tests/test-dataset.py:142–188

Prepare GloVe dataset by converting it to parquet format with chunking

(dataset_dir)

def prepare_glove_dataset(dataset_dir):
    """Prepare GloVe dataset by converting it to parquet format with chunking"""
    zip_path = os.path.join(dataset_dir, "glove.6B.zip")
    if not os.path.exists(zip_path):
        print("Downloading GloVe dataset...")
        download_file(datasets["glove-100"]["url"], zip_path)

    print("Extracting GloVe dataset...")
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extract("glove.6B.100d.txt", dataset_dir)

    print("Converting GloVe to parquet format with chunking...")
    # Read the GloVe text file in chunks
    chunk_size = 200_000
    chunk = []
    file_index = 0
    
    with open(
        os.path.join(dataset_dir, "glove.6B.100d.txt"), "r", encoding="utf-8"
    ) as f:
        for i, line in enumerate(f):
            values = line.split()
            word = values[0]
            vector = [float(x) for x in values[1:]]
            chunk.append({"id": word, "embeddings": vector})
            
            if (i + 1) % chunk_size == 0:
                # Convert chunk to DataFrame and save
                df = pd.DataFrame(chunk)
                parquet_path = os.path.join(dataset_dir, f"test{file_index}.parquet")
                print(f"Saving chunk {file_index} to {parquet_path}")
                df.to_parquet(parquet_path)
                
                # Clear chunk for next batch
                chunk.clear()
                file_index += 1
        
        # Save any remaining data
        if chunk:
            df = pd.DataFrame(chunk)
            parquet_path = os.path.join(dataset_dir, f"test{file_index}.parquet")
            print(f"Saving remaining chunk to {parquet_path}")
            df.to_parquet(parquet_path)

    # Clean up temporary files
    os.remove(os.path.join(dataset_dir, "glove.6B.100d.txt"))
    os.remove(zip_path)


def prepare_dataset(dataset_name):