create_queries.py

"""
This script creates a queries2.csv file from the train2.csv and test2.csv files with the correct `split` column set for each row.
"""


import csv
import sys
import pandas as pd
import argparse

csv.field_size_limit(sys.maxsize)

def create_queries(input_train_file, input_test_file, output_file):
    # Read the train and test splits
    print(f"Reading {input_train_file}...")
    train_df = pd.read_csv(input_train_file)
    print(f"Train rows: {len(train_df)}")

    print(f"Reading {input_test_file}...")
    test_df = pd.read_csv(input_test_file)
    print(f"Test rows: {len(test_df)}")

    # Get unique questions from each split (keep first occurrence to get the id)
    train_unique = train_df.drop_duplicates(subset=['question'], keep='first')[['id', 'question']].copy()
    train_unique['split'] = 'train'
    print(f"Unique train questions: {len(train_unique)}")

    test_unique = test_df.drop_duplicates(subset=['question'], keep='first')[['id', 'question']].copy()
    test_unique['split'] = 'test'
    print(f"Unique test questions: {len(test_unique)}")

    # Combine train and test
    combined = pd.concat([train_unique, test_unique], ignore_index=True)
    print(f"Total unique questions: {len(combined)}")

    # Create the queries dataframe with the required schema
    queries_df = pd.DataFrame({
        'index': range(len(combined)),
        'task': 'none',
        'split': combined['split'],
        'id': combined['id'],
        'question': combined['question']
    })

    # Write to CSV
    print(f"Writing {output_file}...")
    queries_df.to_csv(output_file, index=False)

    print("Done!")
    print(f"Final {output_file} has {len(queries_df)} rows")

def main():
    parser = argparse.ArgumentParser(
        description='Create queries.csv file from train and test sets.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  python create_queries.py input_train.csv input_test.csv queries.csv
        '''
    )
    parser.add_argument('--input-train-file', help='Input train CSV file')
    parser.add_argument('--input-test-file', help='Input test CSV file')
    parser.add_argument('--output-file', help='Output CSV file')
    
    args = parser.parse_args()
    
    create_queries(args.input_train_file, args.input_test_file, args.output_file)

if __name__ == '__main__':
    main()