code_switching_generator.py

"""
Code-Switching Data Generator
Creates synthetic Hinglish, Benglish, and Marathglish data from monolingual datasets
"""

import random
import re
from typing import Dict, List, Tuple, Optional
import json
import os

class CodeSwitchingGenerator:
    """
    Generates synthetic code-switched data by replacing words with English equivalents
    """
    
    def __init__(self, replacement_prob: float = 0.25, seed: int = 42):
        """
        Args:
            replacement_prob: Probability of replacing a word with English
            seed: Random seed for reproducibility
        """
        self.replacement_prob = replacement_prob
        random.seed(seed)
        
        # Common Hindi-English word mappings
        self.hindi_english_map = {
            # Time expressions
            'कल': 'yesterday',
            'आज': 'today',
            'कल': 'tomorrow',
            'अभी': 'now',
            'बाद': 'later',
            'पहले': 'before',
            
            # Common verbs
            'जाना': 'go',
            'आना': 'come',
            'करना': 'do',
            'देखना': 'see',
            'सुनना': 'hear',
            'बोलना': 'speak',
            'खाना': 'eat',
            'पीना': 'drink',
            'सोना': 'sleep',
            'उठना': 'wake up',
            
            # Common nouns
            'घर': 'home',
            'ऑफिस': 'office',
            'स्कूल': 'school',
            'मार्केट': 'market',
            'दुकान': 'shop',
            'रेस्टोरेंट': 'restaurant',
            'हॉस्पिटल': 'hospital',
            'बैंक': 'bank',
            
            # Food items
            'खाना': 'food',
            'पानी': 'water',
            'चाय': 'tea',
            'कॉफी': 'coffee',
            'दूध': 'milk',
            
            # Technology
            'फोन': 'phone',
            'कंप्यूटर': 'computer',
            'इंटरनेट': 'internet',
            'ईमेल': 'email',
            'मैसेज': 'message',
            
            # Actions (common in voice agents)
            'ऑर्डर': 'order',
            'बुक': 'book',
            'कैंसल': 'cancel',
            'चेक': 'check',
            'सर्च': 'search',
            'शेयर': 'share',
            
            # Common adjectives
            'अच्छा': 'good',
            'बुरा': 'bad',
            'बड़ा': 'big',
            'छोटा': 'small',
            'नया': 'new',
            'पुराना': 'old',
            
            # Numbers (often used in English)
            'एक': 'one',
            'दो': 'two',
            'तीन': 'three',
            'चार': 'four',
            'पांच': 'five',
        }
        
        # Common Bengali-English word mappings
        self.bengali_english_map = {
            # Time expressions
            'আজ': 'today',
            'কাল': 'tomorrow',
            'গতকাল': 'yesterday',
            'এখন': 'now',
            'পরে': 'later',
            
            # Common verbs
            'যাওয়া': 'go',
            'আসা': 'come',
            'করা': 'do',
            'দেখা': 'see',
            'শোনা': 'hear',
            'বলা': 'speak',
            'খাওয়া': 'eat',
            
            # Common nouns
            'বাড়ি': 'home',
            'অফিস': 'office',
            'স্কুল': 'school',
            'মার্কেট': 'market',
            'দোকান': 'shop',
            
            # Technology
            'ফোন': 'phone',
            'কম্পিউটার': 'computer',
            'ইন্টারনেট': 'internet',
            'মেসেজ': 'message',
            
            # Actions
            'অর্ডার': 'order',
            'বুক': 'book',
            'ক্যান্সেল': 'cancel',
            'চেক': 'check',
            
            # Common adjectives
            'ভালো': 'good',
            'খারাপ': 'bad',
            'বড়': 'big',
            'ছোট': 'small',
        }
        
        # Common Marathi-English word mappings
        self.marathi_english_map = {
            # Time expressions
            'आज': 'today',
            'उद्या': 'tomorrow',
            'काल': 'yesterday',
            'आता': 'now',
            
            # Common verbs
            'जाणे': 'go',
            'येणे': 'come',
            'करणे': 'do',
            'बघणे': 'see',
            'खाणे': 'eat',
            
            # Common nouns
            'घर': 'home',
            'ऑफिस': 'office',
            'शाळा': 'school',
            'बाजार': 'market',
            
            # Technology
            'फोन': 'phone',
            'संगणक': 'computer',
            'मेसेज': 'message',
            
            # Actions
            'ऑर्डर': 'order',
            'बुक': 'book',
        }
        
        self.language_maps = {
            'hindi': self.hindi_english_map,
            'bengali': self.bengali_english_map,
            'marathi': self.marathi_english_map,
        }
    
    def generate_code_switched_text(
        self, 
        text: str, 
        source_language: str,
        replacement_prob: Optional[float] = None
    ) -> Tuple[str, List[str]]:
        """
        Generate code-switched version of text
        
        Args:
            text: Original text in source language
            source_language: Language of the text ('hindi', 'bengali', 'marathi')
            replacement_prob: Override default replacement probability
            
        Returns:
            Tuple of (code_switched_text, list_of_languages_used)
        """
        if replacement_prob is None:
            replacement_prob = self.replacement_prob
        
        if source_language not in self.language_maps:
            return text, [source_language]
        
        word_map = self.language_maps[source_language]
        words = text.split()
        languages_used = [source_language]
        
        new_words = []
        for word in words:
            # Clean punctuation for matching
            clean_word = re.sub(r'[।,!?;:।]', '', word)
            
            # Check if word can be replaced
            if clean_word in word_map and random.random() < replacement_prob:
                # Replace with English equivalent
                english_word = word_map[clean_word]
                new_words.append(english_word)
                if 'english' not in languages_used:
                    languages_used.append('english')
            else:
                new_words.append(word)
        
        return ' '.join(new_words), languages_used
    
    def generate_conversational_patterns(
        self,
        source_language: str
    ) -> List[Tuple[str, List[str]]]:
        """
        Generate common conversational code-switching patterns
        
        Args:
            source_language: Base language for patterns
            
        Returns:
            List of (text, languages) tuples
        """
        patterns = []
        
        if source_language == 'hindi':
            patterns = [
                # Command patterns
                ("मुझे pizza order करना है", ['hindi', 'english']),
                ("please मेरी help करो", ['hindi', 'english']),
                ("I want to market जाना है", ['hindi', 'english']),
                ("can you check करो मेरा booking", ['hindi', 'english']),
                
                # Question patterns
                ("क्या you can help me", ['hindi', 'english']),
                ("यह item available है क्या", ['hindi', 'english']),
                ("कब है मेरा appointment", ['hindi', 'english']),
                
                # Confirmation patterns
                ("हां yes that's correct", ['hindi', 'english']),
                ("no मुझे वो नहीं चाहिए", ['hindi', 'english']),
                ("okay ठीक है", ['hindi', 'english']),
                
                # Mixed sentences
                ("main yesterday market गया था", ['hindi', 'english']),
                ("मैं अभी office में हूं", ['hindi', 'english']),
                ("please call करो later", ['hindi', 'english']),
            ]
        
        elif source_language == 'bengali':
            patterns = [
                # Command patterns
                ("আমি pizza order করতে চাই", ['bengali', 'english']),
                ("please আমার help করো", ['bengali', 'english']),
                ("I want to market যেতে চাই", ['bengali', 'english']),
                
                # Question patterns
                ("এটা available আছে কি", ['bengali', 'english']),
                ("কখন আছে আমার appointment", ['bengali', 'english']),
                
                # Confirmation patterns
                ("হ্যাঁ yes that's correct", ['bengali', 'english']),
                ("no আমি ওটা চাই না", ['bengali', 'english']),
                
                # Mixed sentences
                ("আমি yesterday market গিয়েছিলাম", ['bengali', 'english']),
                ("আমি এখন office এ আছি", ['bengali', 'english']),
            ]
        
        elif source_language == 'marathi':
            patterns = [
                # Command patterns
                ("मला pizza order करायचा आहे", ['marathi', 'english']),
                ("please माझी help करा", ['marathi', 'english']),
                
                # Question patterns
                ("हे available आहे का", ['marathi', 'english']),
                ("केव्हा आहे माझी appointment", ['marathi', 'english']),
                
                # Confirmation patterns
                ("होय yes that's correct", ['marathi', 'english']),
                ("no मला ते नको", ['marathi', 'english']),
                
                # Mixed sentences
                ("मी yesterday market गेलो होतो", ['marathi', 'english']),
            ]
        
        return patterns
    
    def augment_dataset(
        self,
        samples: List[Dict],
        source_language: str,
        augmentation_factor: int = 2,
        include_conversational: bool = True
    ) -> List[Dict]:
        """
        Create augmented dataset with code-switching
        
        Args:
            samples: List of dicts with 'audio_path' and 'text' keys
            source_language: Language of samples
            augmentation_factor: How many CS versions to create per sample
            include_conversational: Whether to add conversational patterns
            
        Returns:
            Augmented list of samples
        """
        augmented = []
        
        # Add original samples
        for sample in samples:
            augmented.append({
                **sample,
                'language': source_language,
                'is_code_switched': False
            })
        
        # Create code-switched versions
        for sample in samples:
            for i in range(augmentation_factor):
                # Vary replacement probability
                prob = self.replacement_prob + random.uniform(-0.1, 0.1)
                prob = max(0.1, min(0.4, prob))  # Clamp between 0.1 and 0.4
                
                cs_text, languages = self.generate_code_switched_text(
                    sample['text'],
                    source_language,
                    replacement_prob=prob
                )
                
                # Only add if actually code-switched
                if len(languages) > 1:
                    augmented.append({
                        'audio_path': sample['audio_path'],
                        'text': cs_text,
                        'language': f"{source_language}lish",  # e.g., 'hinglish'
                        'is_code_switched': True,
                        'languages': languages
                    })
        
        # Add conversational patterns (synthetic text, would need TTS for audio)
        if include_conversational:
            patterns = self.generate_conversational_patterns(source_language)
            for text, languages in patterns:
                augmented.append({
                    'audio_path': None,  # Would need TTS or real recording
                    'text': text,
                    'language': f"{source_language}lish",
                    'is_code_switched': True,
                    'languages': languages,
                    'is_synthetic_text': True
                })
        
        return augmented
    
    def save_augmented_manifest(
        self,
        augmented_samples: List[Dict],
        output_path: str
    ):
        """Save augmented dataset to JSON manifest"""
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(augmented_samples, f, ensure_ascii=False, indent=2)
        print(f"Saved {len(augmented_samples)} samples to {output_path}")


def demo():
    """Demonstration of code-switching generation"""
    generator = CodeSwitchingGenerator(replacement_prob=0.3)
    
    # Test Hindi
    hindi_text = "मैं कल मार्केट जाऊंगा और खाना खरीदूंगा"
    cs_text, langs = generator.generate_code_switched_text(hindi_text, 'hindi')
    print(f"Original: {hindi_text}")
    print(f"Code-switched: {cs_text}")
    print(f"Languages: {langs}")
    print()
    
    # Test Bengali
    bengali_text = "আমি কাল মার্কেট যাব এবং খাবার কিনব"
    cs_text, langs = generator.generate_code_switched_text(bengali_text, 'bengali')
    print(f"Original: {bengali_text}")
    print(f"Code-switched: {cs_text}")
    print(f"Languages: {langs}")
    print()
    
    # Show conversational patterns
    print("Conversational patterns (Hindi):")
    for text, langs in generator.generate_conversational_patterns('hindi')[:5]:
        print(f"  {text} - {langs}")


if __name__ == "__main__":
    demo()