langtune 0.1.19__py3-none-any.whl

langtune/py.typed

@@ -0,0 +1,2 @@
+# Marker file for PEP 561.
+# The langtune package uses inline type annotations.

langtune/schedulers.py

@@ -0,0 +1,234 @@
+"""
+schedulers.py: Learning rate schedulers for Langtune
+
+Provides various learning rate scheduling strategies.
+"""
+
+import math
+import torch
+from torch.optim.lr_scheduler import _LRScheduler
+from typing import Optional, List
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class WarmupScheduler(_LRScheduler):
+    """Linear warmup scheduler."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_steps: int,
+        base_lr: float = None,
+        last_epoch: int = -1
+    ):
+        self.warmup_steps = warmup_steps
+        self.base_lr = base_lr
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            warmup_factor = (self.last_epoch + 1) / max(self.warmup_steps, 1)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+        return self.base_lrs
+
+
+class CosineAnnealingWithWarmup(_LRScheduler):
+    """Cosine annealing with linear warmup."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        total_steps: int,
+        warmup_steps: int = 0,
+        min_lr: float = 0.0,
+        last_epoch: int = -1
+    ):
+        self.total_steps = total_steps
+        self.warmup_steps = warmup_steps
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            # Linear warmup
+            warmup_factor = (self.last_epoch + 1) / max(self.warmup_steps, 1)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+        else:
+            # Cosine annealing
+            progress = (self.last_epoch - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
+            return [
+                self.min_lr + 0.5 * (base_lr - self.min_lr) * (1 + math.cos(math.pi * progress))
+                for base_lr in self.base_lrs
+            ]
+
+
+class LinearDecayWithWarmup(_LRScheduler):
+    """Linear decay with warmup."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        total_steps: int,
+        warmup_steps: int = 0,
+        min_lr: float = 0.0,
+        last_epoch: int = -1
+    ):
+        self.total_steps = total_steps
+        self.warmup_steps = warmup_steps
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            # Linear warmup
+            warmup_factor = (self.last_epoch + 1) / max(self.warmup_steps, 1)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+        else:
+            # Linear decay
+            progress = (self.last_epoch - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
+            return [
+                max(self.min_lr, base_lr * (1 - progress))
+                for base_lr in self.base_lrs
+            ]
+
+
+class PolynomialDecayWithWarmup(_LRScheduler):
+    """Polynomial decay with warmup."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        total_steps: int,
+        warmup_steps: int = 0,
+        power: float = 2.0,
+        min_lr: float = 0.0,
+        last_epoch: int = -1
+    ):
+        self.total_steps = total_steps
+        self.warmup_steps = warmup_steps
+        self.power = power
+        self.min_lr = min_lr
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            warmup_factor = (self.last_epoch + 1) / max(self.warmup_steps, 1)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+        else:
+            progress = (self.last_epoch - self.warmup_steps) / max(self.total_steps - self.warmup_steps, 1)
+            decay_factor = (1 - progress) ** self.power
+            return [
+                max(self.min_lr, base_lr * decay_factor)
+                for base_lr in self.base_lrs
+            ]
+
+
+class ConstantWithWarmup(_LRScheduler):
+    """Constant learning rate with warmup."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        warmup_steps: int = 0,
+        last_epoch: int = -1
+    ):
+        self.warmup_steps = warmup_steps
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            warmup_factor = (self.last_epoch + 1) / max(self.warmup_steps, 1)
+            return [base_lr * warmup_factor for base_lr in self.base_lrs]
+        return self.base_lrs
+
+
+class OneCycleLRWithWarmup(_LRScheduler):
+    """1cycle LR policy with customizable warmup."""
+    
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        max_lr: float,
+        total_steps: int,
+        pct_start: float = 0.3,
+        div_factor: float = 25.0,
+        final_div_factor: float = 10000.0,
+        last_epoch: int = -1
+    ):
+        self.max_lr = max_lr
+        self.total_steps = total_steps
+        self.pct_start = pct_start
+        self.div_factor = div_factor
+        self.final_div_factor = final_div_factor
+        
+        # Calculate phase boundaries
+        self.warmup_steps = int(total_steps * pct_start)
+        self.cooldown_steps = total_steps - self.warmup_steps
+        
+        super().__init__(optimizer, last_epoch)
+    
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            # Warmup phase: linear increase
+            progress = self.last_epoch / max(self.warmup_steps, 1)
+            start_lr = self.max_lr / self.div_factor
+            return [start_lr + (self.max_lr - start_lr) * progress for _ in self.base_lrs]
+        else:
+            # Cooldown phase: cosine decay
+            progress = (self.last_epoch - self.warmup_steps) / max(self.cooldown_steps, 1)
+            end_lr = self.max_lr / self.final_div_factor
+            return [
+                end_lr + 0.5 * (self.max_lr - end_lr) * (1 + math.cos(math.pi * progress))
+                for _ in self.base_lrs
+            ]
+
+
+def get_scheduler(
+    name: str,
+    optimizer: torch.optim.Optimizer,
+    total_steps: int,
+    warmup_steps: int = 0,
+    **kwargs
+) -> _LRScheduler:
+    """
+    Get a scheduler by name.
+    
+    Args:
+        name: Scheduler name ('cosine', 'linear', 'constant', 'polynomial', 'onecycle')
+        optimizer: Optimizer
+        total_steps: Total training steps
+        warmup_steps: Warmup steps
+        **kwargs: Additional scheduler arguments
+        
+    Returns:
+        Learning rate scheduler
+    """
+    schedulers = {
+        "cosine": CosineAnnealingWithWarmup,
+        "linear": LinearDecayWithWarmup,
+        "constant": ConstantWithWarmup,
+        "polynomial": PolynomialDecayWithWarmup,
+        "warmup": WarmupScheduler,
+    }
+    
+    if name == "onecycle":
+        max_lr = kwargs.get("max_lr", optimizer.param_groups[0]['lr'])
+        return OneCycleLRWithWarmup(
+            optimizer, max_lr, total_steps,
+            pct_start=warmup_steps / total_steps if total_steps > 0 else 0.3,
+            **{k: v for k, v in kwargs.items() if k != "max_lr"}
+        )
+    
+    if name not in schedulers:
+        raise ValueError(f"Unknown scheduler: {name}. Options: {list(schedulers.keys())}")
+    
+    scheduler_cls = schedulers[name]
+    
+    if name == "warmup":
+        return scheduler_cls(optimizer, warmup_steps, **kwargs)
+    elif name == "constant":
+        return scheduler_cls(optimizer, warmup_steps, **kwargs)
+    else:
+        return scheduler_cls(optimizer, total_steps, warmup_steps, **kwargs)

langtune/tokenizers.py

@@ -0,0 +1,275 @@
+"""
+tokenizers.py: Tokenization utilities for Langtune
+
+Provides tokenization helpers and wrappers for various tokenizers.
+"""
+
+import re
+from typing import List, Dict, Optional, Union
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class CharacterTokenizer:
+    """Simple character-level tokenizer."""
+    
+    def __init__(self, vocab_size: int = 256):
+        self.vocab_size = vocab_size
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+    
+    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        """Encode text to character IDs."""
+        tokens = []
+        if add_special_tokens:
+            tokens.append(self.bos_token_id)
+        
+        for char in text:
+            token_id = ord(char) % (self.vocab_size - 4) + 4
+            tokens.append(token_id)
+        
+        if add_special_tokens:
+            tokens.append(self.eos_token_id)
+        return tokens
+    
+    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
+        """Decode token IDs to text."""
+        chars = []
+        for tid in token_ids:
+            if skip_special_tokens and tid < 4:
+                continue
+            if tid >= 4:
+                chars.append(chr((tid - 4) % 128 + 32))
+        return ''.join(chars)
+    
+    def __call__(self, text: str, **kwargs) -> Dict[str, List[int]]:
+        return {"input_ids": self.encode(text, **kwargs)}
+
+
+class WordTokenizer:
+    """Simple word-level tokenizer with vocabulary."""
+    
+    def __init__(self, vocab: Optional[Dict[str, int]] = None, max_vocab_size: int = 32000):
+        self.max_vocab_size = max_vocab_size
+        self.vocab = vocab or {}
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        
+        # Special tokens
+        self.pad_token = "<pad>"
+        self.unk_token = "<unk>"
+        self.bos_token = "<s>"
+        self.eos_token = "</s>"
+        
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+        
+        if not self.vocab:
+            self.vocab = {
+                self.pad_token: 0,
+                self.unk_token: 1,
+                self.bos_token: 2,
+                self.eos_token: 3
+            }
+            self.inv_vocab = {v: k for k, v in self.vocab.items()}
+    
+    def fit(self, texts: List[str], min_freq: int = 1):
+        """Build vocabulary from texts."""
+        word_counts = {}
+        for text in texts:
+            for word in self._tokenize(text):
+                word_counts[word] = word_counts.get(word, 0) + 1
+        
+        # Sort by frequency
+        sorted_words = sorted(word_counts.items(), key=lambda x: -x[1])
+        
+        # Add to vocabulary
+        for word, count in sorted_words:
+            if count < min_freq:
+                break
+            if len(self.vocab) >= self.max_vocab_size:
+                break
+            if word not in self.vocab:
+                self.vocab[word] = len(self.vocab)
+        
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        logger.info(f"Built vocabulary with {len(self.vocab)} tokens")
+    
+    def _tokenize(self, text: str) -> List[str]:
+        """Split text into words."""
+        text = text.lower()
+        return re.findall(r'\b\w+\b|[^\w\s]', text)
+    
+    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        """Encode text to token IDs."""
+        tokens = []
+        if add_special_tokens:
+            tokens.append(self.bos_token_id)
+        
+        for word in self._tokenize(text):
+            tokens.append(self.vocab.get(word, self.unk_token_id))
+        
+        if add_special_tokens:
+            tokens.append(self.eos_token_id)
+        return tokens
+    
+    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
+        """Decode token IDs to text."""
+        words = []
+        for tid in token_ids:
+            if skip_special_tokens and tid < 4:
+                continue
+            word = self.inv_vocab.get(tid, self.unk_token)
+            if not skip_special_tokens or not word.startswith("<"):
+                words.append(word)
+        return ' '.join(words)
+    
+    def __call__(self, text: str, **kwargs) -> Dict[str, List[int]]:
+        return {"input_ids": self.encode(text, **kwargs)}
+    
+    def save(self, path: str):
+        """Save vocabulary to file."""
+        import json
+        with open(path, 'w') as f:
+            json.dump(self.vocab, f, indent=2)
+    
+    @classmethod
+    def load(cls, path: str) -> "WordTokenizer":
+        """Load vocabulary from file."""
+        import json
+        with open(path) as f:
+            vocab = json.load(f)
+        return cls(vocab=vocab)
+
+
+class BPETokenizer:
+    """Simple Byte-Pair Encoding tokenizer."""
+    
+    def __init__(self, vocab_size: int = 8000):
+        self.vocab_size = vocab_size
+        self.merges = {}
+        self.vocab = {}
+        
+        self.pad_token_id = 0
+        self.unk_token_id = 1
+        self.bos_token_id = 2
+        self.eos_token_id = 3
+    
+    def _get_pairs(self, word: List[str]) -> Dict[tuple, int]:
+        """Get pairs of consecutive symbols."""
+        pairs = {}
+        for i in range(len(word) - 1):
+            pair = (word[i], word[i + 1])
+            pairs[pair] = pairs.get(pair, 0) + 1
+        return pairs
+    
+    def fit(self, texts: List[str], num_merges: int = None):
+        """Learn BPE merges from texts."""
+        num_merges = num_merges or (self.vocab_size - 256)
+        
+        # Initialize vocabulary with characters
+        word_freqs = {}
+        for text in texts:
+            for word in text.split():
+                word = ' '.join(list(word)) + ' </w>'
+                word_freqs[word] = word_freqs.get(word, 0) + 1
+        
+        # Build initial vocab
+        self.vocab = {chr(i): i for i in range(256)}
+        self.vocab['</w>'] = 256
+        
+        # Learn merges
+        for i in range(num_merges):
+            pairs = {}
+            for word, freq in word_freqs.items():
+                word_pairs = self._get_pairs(word.split())
+                for pair, count in word_pairs.items():
+                    pairs[pair] = pairs.get(pair, 0) + count * freq
+            
+            if not pairs:
+                break
+            
+            best_pair = max(pairs, key=pairs.get)
+            new_token = ''.join(best_pair)
+            
+            if len(self.vocab) >= self.vocab_size:
+                break
+            
+            self.merges[best_pair] = len(self.vocab)
+            self.vocab[new_token] = len(self.vocab)
+            
+            # Update word_freqs
+            new_word_freqs = {}
+            pattern = ' '.join(best_pair)
+            replacement = new_token
+            for word, freq in word_freqs.items():
+                new_word = word.replace(pattern, replacement)
+                new_word_freqs[new_word] = freq
+            word_freqs = new_word_freqs
+        
+        logger.info(f"Learned {len(self.merges)} BPE merges")
+    
+    def encode(self, text: str, add_special_tokens: bool = True) -> List[int]:
+        """Encode text using BPE."""
+        tokens = []
+        if add_special_tokens:
+            tokens.append(self.bos_token_id)
+        
+        for word in text.split():
+            word = ' '.join(list(word)) + ' </w>'
+            while True:
+                pairs = self._get_pairs(word.split())
+                if not pairs:
+                    break
+                
+                # Find best pair that exists in merges
+                best_pair = None
+                best_rank = float('inf')
+                for pair in pairs:
+                    if pair in self.merges and self.merges[pair] < best_rank:
+                        best_pair = pair
+                        best_rank = self.merges[pair]
+                
+                if best_pair is None:
+                    break
+                
+                pattern = ' '.join(best_pair)
+                replacement = ''.join(best_pair)
+                word = word.replace(pattern, replacement)
+            
+            for token in word.split():
+                tokens.append(self.vocab.get(token, self.unk_token_id))
+        
+        if add_special_tokens:
+            tokens.append(self.eos_token_id)
+        return tokens
+    
+    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
+        """Decode BPE tokens."""
+        inv_vocab = {v: k for k, v in self.vocab.items()}
+        tokens = []
+        for tid in token_ids:
+            if skip_special_tokens and tid < 4:
+                continue
+            token = inv_vocab.get(tid, '')
+            tokens.append(token.replace('</w>', ' '))
+        return ''.join(tokens).strip()
+
+
+def get_tokenizer(name: str = "character", **kwargs):
+    """Get a tokenizer by name."""
+    tokenizers = {
+        "character": CharacterTokenizer,
+        "char": CharacterTokenizer,
+        "word": WordTokenizer,
+        "bpe": BPETokenizer
+    }
+    
+    if name not in tokenizers:
+        raise ValueError(f"Unknown tokenizer: {name}. Options: {list(tokenizers.keys())}")
+    
+    return tokenizers[name](**kwargs)