code/data/indirect_idx/generate.py

import random
import string
from typing import List, Tuple, Optional

class IndirectIdxDatasetGenerator:
    def __init__(self, max_shift: int = 10, min_length: int = 20, max_length: int = 40):
        self.max_shift = max_shift
        self.min_length = min_length
        self.max_length = max_length
        self.all_chars = list(string.ascii_uppercase + string.ascii_lowercase)
    
    def generate_sequence(self) -> List[str]:
        sequence_length = random.randint(self.min_length, self.max_length)
        sampled_chars = random.sample(self.all_chars, sequence_length)
        random.shuffle(sampled_chars)
        return sampled_chars
    
    def is_valid_shift(self, sequence: List[str], source_char: str, shift: int) -> bool:
        """ Check if a shift is valid (target position exists and no wrap-around). """
        try:
            source_idx = sequence.index(source_char)
            target_idx = source_idx + shift
            return 0 <= target_idx < len(sequence)
        except ValueError:
            return False
    
    def get_target_char(self, sequence: List[str], source_char: str, shift: int) -> Optional[str]:
        if not self.is_valid_shift(sequence, source_char, shift):
            return None
        
        source_idx = sequence.index(source_char)
        target_idx = source_idx + shift
        return sequence[target_idx]
    
    def generate_sample(self) -> Tuple[str, str]:
        max_attempts = 1000
        attempts = 0
        
        while attempts < max_attempts:
            sequence = self.generate_sequence()            
            source_char = random.choice(sequence)
            shift = random.randint(-self.max_shift, self.max_shift)
            target_char = self.get_target_char(sequence, source_char, shift)
            if target_char is not None:
                seq_str = ''.join(sequence)
                input_str = f"{seq_str}, {source_char}, {shift:+d}"
                return input_str, target_char
            
            attempts += 1
        
        raise RuntimeError("Could not generate valid sample after maximum attempts")
    
    def generate_dataset(self, num_samples: int) -> List[Tuple[str, str]]:
        dataset = []
        for _ in range(num_samples):
            sample = self.generate_sample()
            dataset.append(sample)
        return dataset
    
    def save_dataset(self, dataset: List[Tuple[str, str]], filename: str):
        with open(filename, 'w', encoding='utf-8') as f:
            for input_str, target in dataset:
                f.write(f"{input_str}, {target}\n")
        
def main():
    """Example usage of the dataset generator."""
    # Initialize generator with custom parameters
    shift, min_len, max_len = 15, 20, 40
    generator = IndirectIdxDatasetGenerator(max_shift=shift, min_length=min_len, max_length=max_len)
    
    # Generate a small sample dataset
    print("Generating sample dataset...")
    dataset = generator.generate_dataset(10)
    
    # Display samples
    print("\nSample data:")
    print("-" * 80)
    for input_str, target in dataset:
        print(f"{input_str}, {target}")
    
    # Generate larger dataset and save to files
    print(f"\nGenerating larger dataset...")
    large_dataset = generator.generate_dataset(1020000)
    
    # Save in different formats
    fname = f"data/indirect_idx/ds_minl{min_len}_maxl{max_len}_shift_{shift}.txt"
    generator.save_dataset(large_dataset, fname)
    print(f"Saving dataset to disk as txt file.")
    
    # Show comprehensive statistics
    shifts = [int(sample[0].split(', ')[2]) for sample in large_dataset]
    lengths = [len(sample[0].split(', ')[0]) for sample in large_dataset]
    
    print(f"\nDataset statistics:")
    print(f"Total samples: {len(large_dataset)}")
    print(f"Shift range: {min(shifts)} to {max(shifts)}")
    print(f"Average shift: {sum(shifts)/len(shifts):.2f}")
    print(f"Sequence length range: {min(lengths)} to {max(lengths)}")
    print(f"Average sequence length: {sum(lengths)/len(lengths):.1f}")
    
    # Show character distribution
    all_chars_used = set()
    for sample in large_dataset:
        sequence = sample[0].split(', ')[0]
        all_chars_used.update(sequence)
    
    uppercase_count = sum(1 for c in all_chars_used if c.isupper())
    lowercase_count = sum(1 for c in all_chars_used if c.islower())
    
    print(f"Character usage:")
    print(f"Unique uppercase letters used: {uppercase_count}/26")
    print(f"Unique lowercase letters used: {lowercase_count}/26")
    print(f"Total unique characters used: {len(all_chars_used)}/52")

if __name__ == "__main__":
    main()