Fix API model loading: Copy src directory and update Dockerfile

Dockerfile

@@ -8,6 +8,7 @@ RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy application
 COPY api.py .
+COPY src/ src/
 COPY model/ model/
 
 # Expose port

api.py

@@ -81,39 +81,8 @@ class LSTMLanguageModel(nn.Module):
 # =============================================================================
 # Trigram Model
 # =============================================================================
-class TrigramLM:
-    def __init__(self, smoothing: float = 1.0):
-        self.smoothing = smoothing
-        self.unigram_counts = {}
-        self.bigram_counts = {}
-        self.trigram_counts = {}
-        self.vocab = set()
-    
-    def probability(self, w3: str, w1: str, w2: str) -> float:
-        trigram_count = self.trigram_counts.get((w1, w2, w3), 0)
-        bigram_count = self.bigram_counts.get((w1, w2), 0)
-        vocab_size = len(self.vocab)
-        numerator = trigram_count + self.smoothing
-        denominator = bigram_count + (self.smoothing * vocab_size)
-        return numerator / denominator if denominator > 0 else 0.0
-    
-    def predict_next_words(self, context: str, top_k: int = 5) -> List[Tuple[str, float]]:
-        words = context.lower().split()
-        if len(words) == 0:
-            w1, w2 = START_TOKEN, START_TOKEN
-        elif len(words) == 1:
-            w1, w2 = START_TOKEN, words[0]
-        else:
-            w1, w2 = words[-2], words[-1]
-        
-        candidates = []
-        for word in self.vocab:
-            if word not in (START_TOKEN, END_TOKEN, '<s>', '</s>'):
-                prob = self.probability(word, w1, w2)
-                candidates.append((word, prob))
-        
-        candidates.sort(key=lambda x: x[1], reverse=True)
-        return candidates[:top_k]
+# Import directly from src to ensure compatibility with pickle
+from src.trigram_model import TrigramLM
 
 # =============================================================================
 # Global Models (loaded once at startup)
@@ -123,16 +92,6 @@ word_to_idx = None
 idx_to_word = None
 trigram_model = None
 
-# =============================================================================
-# Custom Unpickler to fix the 'src' module error
-# =============================================================================
-class PatchingUnpickler(pickle.Unpickler):
-    def find_class(self, module, name):
-        # If the pickle creates a dependency on 'src', redirect it to __main__
-        if module.startswith("src") and name == "TrigramLM":
-            return TrigramLM
-        return super().find_class(module, name)
-
 @app.on_event("startup")
 async def load_models():
     global lstm_model, word_to_idx, idx_to_word, trigram_model
@@ -151,11 +110,10 @@ async def load_models():
     except Exception as e:
         print(f"Failed to load LSTM model: {e}")
 
-    # 2. Load Trigram (Using the Custom Unpickler)
+    # 2. Load Trigram
     try:
         with open('model/trigram_model.pkl', 'rb') as f:
-            # Use PatchingUnpickler instead of standard pickle.load
-            trigram_model = PatchingUnpickler(f).load()
+            trigram_model = pickle.load(f)
         print(f"Trigram model loaded! Vocab size: {len(trigram_model.vocab)}")
     except Exception as e:
         print(f"Failed to load Trigram model: {e}")

src/__init__.py

@@ -0,0 +1,2 @@
+# Nigerian English/Pidgin Next-Word Prediction
+# Trigram Language Model Baseline

src/data_loader.py

@@ -0,0 +1,148 @@
+"""
+Data loading utilities for NaijaSenti and BBC Pidgin datasets.
+
+Loads Nigerian Pidgin text from multiple sources for language modeling.
+Sentiment/category labels are ignored.
+"""
+
+from datasets import load_dataset
+from typing import List, Dict, Any, Optional
+import csv
+import os
+
+# Path to BBC Pidgin corpus (relative to project root)
+BBC_PIDGIN_CORPUS_PATH = "bbc_pidgin_scraper/data/pidgin_corpus.csv"
+
+
+def load_naijasenti_pcm() -> Dict[str, List[str]]:
+    """
+    Load the NaijaSenti PCM (Nigerian Pidgin) dataset.
+    
+    Returns:
+        Dict with keys 'train', 'test', 'validation' containing text lists.
+    """
+    dataset = load_dataset("mteb/NaijaSenti", "pcm")
+    
+    result = {}
+    for split in dataset.keys():
+        # Extract text field, ignore sentiment labels
+        result[split] = [example['text'] for example in dataset[split]]
+    
+    return result
+
+
+def load_bbc_pidgin(limit: Optional[int] = None, project_root: Optional[str] = None) -> List[str]:
+    """
+    Load BBC Pidgin articles from the scraped corpus.
+    
+    The corpus contains headlines and article texts scraped from BBC Pidgin.
+    We concatenate headline + text for each article.
+    
+    Args:
+        limit: Maximum number of articles to load. None for all.
+        project_root: Path to project root. Defaults to current working directory.
+        
+    Returns:
+        List of article texts (headline + body combined).
+    """
+    if project_root is None:
+        project_root = os.getcwd()
+    
+    corpus_path = os.path.join(project_root, BBC_PIDGIN_CORPUS_PATH)
+    
+    if not os.path.exists(corpus_path):
+        print(f"Warning: BBC Pidgin corpus not found at {corpus_path}")
+        return []
+    
+    texts = []
+    try:
+        with open(corpus_path, 'r', encoding='utf-8') as f:
+            reader = csv.DictReader(f)
+            for i, row in enumerate(reader):
+                if limit and i >= limit:
+                    break
+                # Combine headline and text
+                headline = row.get('headline', '').strip()
+                text = row.get('text', '').strip()
+                if headline and text:
+                    combined = f"{headline}. {text}"
+                    texts.append(combined)
+                elif text:
+                    texts.append(text)
+    except Exception as e:
+        print(f"Error loading BBC Pidgin corpus: {e}")
+        return []
+    
+    print(f"Loaded {len(texts):,} BBC Pidgin articles")
+    return texts
+
+
+def load_all_texts(include_bbc: bool = True, bbc_limit: Optional[int] = None) -> List[str]:
+    """
+    Load all text from all sources combined.
+    
+    Combines NaijaSenti PCM dataset with BBC Pidgin articles
+    for maximum data coverage.
+    
+    Args:
+        include_bbc: Whether to include BBC Pidgin articles.
+        bbc_limit: Maximum number of BBC articles to include.
+        
+    Returns:
+        List of all text strings from all sources.
+    """
+    all_texts = []
+    
+    # Load NaijaSenti
+    print("Loading NaijaSenti PCM dataset...")
+    splits = load_naijasenti_pcm()
+    for split_name, texts in splits.items():
+        all_texts.extend(texts)
+        print(f"  Loaded {len(texts):,} texts from {split_name} split")
+    
+    naija_total = len(all_texts)
+    print(f"  NaijaSenti total: {naija_total:,} texts")
+    
+    # Load BBC Pidgin
+    if include_bbc:
+        print(f"\nLoading BBC Pidgin corpus (limit={bbc_limit})...")
+        bbc_texts = load_bbc_pidgin(limit=bbc_limit)
+        all_texts.extend(bbc_texts)
+    
+    print(f"\nCombined total: {len(all_texts):,} texts")
+    return all_texts
+
+
+def get_dataset_stats(texts: List[str]) -> Dict[str, Any]:
+    """
+    Compute basic statistics about the dataset.
+    
+    Args:
+        texts: List of text strings.
+        
+    Returns:
+        Dictionary of statistics.
+    """
+    total_chars = sum(len(t) for t in texts)
+    total_words = sum(len(t.split()) for t in texts)
+    
+    return {
+        'num_texts': len(texts),
+        'total_characters': total_chars,
+        'total_words': total_words,
+        'avg_words_per_text': total_words / len(texts) if texts else 0,
+        'avg_chars_per_text': total_chars / len(texts) if texts else 0,
+    }
+
+
+if __name__ == "__main__":
+    # Quick test
+    texts = load_all_texts(include_bbc=True)  # Loads all BBC articles by default
+    stats = get_dataset_stats(texts)
+    print("\nDataset Statistics:")
+    for key, value in stats.items():
+        if isinstance(value, float):
+            print(f"  {key}: {value:.2f}")
+        else:
+            print(f"  {key}: {value:,}")
+

src/preprocessing.py

@@ -0,0 +1,174 @@
+"""
+Text preprocessing pipeline for Nigerian English/Pidgin.
+
+Design principles:
+- Preserve linguistic features of Nigerian Pidgin (slang, contractions, code-switching)
+- Remove noise (URLs, usernames) that don't contribute to language modeling
+- Minimal normalization to avoid losing dialectal patterns
+"""
+
+import re
+from typing import List
+
+
+# Special tokens for sentence boundaries
+START_TOKEN = "<s>"
+END_TOKEN = "</s>"
+
+
+def clean_text(text: str) -> str:
+    """
+    Clean text while preserving Nigerian Pidgin features.
+    
+    Operations:
+    1. Lowercase (case doesn't matter for prediction)
+    2. Remove URLs
+    3. Remove @usernames (Twitter-style)
+    4. Normalize whitespace
+    
+    Preserved:
+    - Contractions (don't, I'm, na'm)
+    - Slang (abi, sha, sef)
+    - Code-switching patterns
+    - Pidgin grammar structures
+    
+    Args:
+        text: Raw text string.
+        
+    Returns:
+        Cleaned text string.
+    """
+    # Lowercase
+    text = text.lower()
+    
+    # Remove URLs
+    text = re.sub(r'https?://\S+', '', text)
+    text = re.sub(r'www\.\S+', '', text)
+    
+    # Remove @usernames
+    text = re.sub(r'@\w+', '', text)
+    
+    # Remove hashtags but keep the word
+    text = re.sub(r'#(\w+)', r'\1', text)
+    
+    # Normalize whitespace
+    text = re.sub(r'\s+', ' ', text)
+    text = text.strip()
+    
+    return text
+
+
+def tokenize(text: str) -> List[str]:
+    """
+    Word-level tokenization for Nigerian Pidgin.
+    
+    Handles:
+    - Standard word boundaries
+    - Punctuation as separate tokens
+    - Preserves contractions as single tokens
+    
+    Args:
+        text: Cleaned text string.
+        
+    Returns:
+        List of tokens.
+    """
+    # Split on whitespace first
+    words = text.split()
+    
+    tokens = []
+    for word in words:
+        # Handle punctuation attached to words
+        # Keep contractions together (don't, I'm, etc.)
+        
+        # Strip leading punctuation
+        while word and word[0] in '.,!?;:"\'-([{':
+            if word[0] not in "'":  # Keep leading apostrophe for contractions
+                tokens.append(word[0])
+            word = word[1:]
+        
+        # Strip trailing punctuation
+        trailing = []
+        while word and word[-1] in '.,!?;:"\'-)]}"':
+            if word[-1] not in "'":  # Keep trailing apostrophe for contractions
+                trailing.insert(0, word[-1])
+            word = word[:-1]
+        
+        if word:
+            tokens.append(word)
+        
+        tokens.extend(trailing)
+    
+    return tokens
+
+
+def preprocess_text(text: str) -> List[str]:
+    """
+    Full preprocessing pipeline: clean + tokenize.
+    
+    Args:
+        text: Raw text string.
+        
+    Returns:
+        List of tokens.
+    """
+    cleaned = clean_text(text)
+    tokens = tokenize(cleaned)
+    return tokens
+
+
+def add_sentence_markers(tokens: List[str]) -> List[str]:
+    """
+    Add start/end markers for sentence boundary modeling.
+    
+    For trigram models, we need context at sentence boundaries.
+    We add two start tokens to provide full context for the first word.
+    
+    Args:
+        tokens: List of tokens from a sentence.
+        
+    Returns:
+        Tokens with boundary markers.
+    """
+    if not tokens:
+        return []
+    return [START_TOKEN, START_TOKEN] + tokens + [END_TOKEN]
+
+
+def preprocess_corpus(texts: List[str]) -> List[List[str]]:
+    """
+    Preprocess entire corpus for language model training.
+    
+    Args:
+        texts: List of raw text strings.
+        
+    Returns:
+        List of tokenized sentences with boundary markers.
+    """
+    processed = []
+    for text in texts:
+        tokens = preprocess_text(text)
+        if tokens:  # Skip empty results
+            marked = add_sentence_markers(tokens)
+            processed.append(marked)
+    return processed
+
+
+if __name__ == "__main__":
+    # Test preprocessing on Nigerian Pidgin examples
+    test_texts = [
+        "I dey go market, you wan follow?",
+        "That guy na correct person sha @handle https://example.com",
+        "Wetin you dey do? Abi you no sabi?",
+        "E don happen before, no be today matter",
+        "How far? Everything dey go well?",
+    ]
+    
+    print("Preprocessing Examples:\n")
+    for text in test_texts:
+        tokens = preprocess_text(text)
+        marked = add_sentence_markers(tokens)
+        print(f"Original: {text}")
+        print(f"Tokens:   {tokens}")
+        print(f"Marked:   {marked}")
+        print()

src/trigram_model.py

@@ -0,0 +1,276 @@
+"""
+Trigram Language Model for Next-Word Prediction.
+
+Implements a statistical trigram model with Laplace (add-one) smoothing
+for Nigerian English/Pidgin next-word prediction.
+
+Mathematical Foundation:
+    P(w_n | w_{n-2}, w_{n-1}) = (C(w_{n-2}, w_{n-1}, w_n) + α) / (C(w_{n-2}, w_{n-1}) + α|V|)
+    
+Where:
+    - C(.) = count of n-gram in training corpus
+    - α = smoothing parameter (1.0 for Laplace)
+    - |V| = vocabulary size
+"""
+
+from collections import Counter
+from typing import List, Tuple, Dict, Optional
+import math
+
+
+class TrigramLM:
+    """
+    Trigram Language Model with Laplace smoothing.
+    
+    Attributes:
+        smoothing: Smoothing parameter (α). Default 1.0 for add-one smoothing.
+        unigram_counts: Counter for single word frequencies.
+        bigram_counts: Counter for word pair frequencies.
+        trigram_counts: Counter for word triple frequencies.
+        vocab: Set of all unique words in training corpus.
+    """
+    
+    def __init__(self, smoothing: float = 1.0):
+        """
+        Initialize the trigram model.
+        
+        Args:
+            smoothing: Laplace smoothing parameter. Higher values provide more
+                      smoothing for unseen n-grams. Default 1.0 (add-one).
+        """
+        self.smoothing = smoothing
+        self.unigram_counts: Counter = Counter()
+        self.bigram_counts: Counter = Counter()
+        self.trigram_counts: Counter = Counter()
+        self.vocab: set = set()
+        self._total_unigrams: int = 0
+    
+    def train(self, sentences: List[List[str]]) -> None:
+        """
+        Train the model by counting n-grams from tokenized sentences.
+        
+        Expects sentences with start/end markers already added:
+        ['<s>', '<s>', 'word1', 'word2', ..., '</s>']
+        
+        Args:
+            sentences: List of tokenized sentences with boundary markers.
+        """
+        for sentence in sentences:
+            # Build vocabulary
+            self.vocab.update(sentence)
+            
+            # Count unigrams
+            for token in sentence:
+                self.unigram_counts[token] += 1
+                self._total_unigrams += 1
+            
+            # Count bigrams
+            for i in range(len(sentence) - 1):
+                bigram = (sentence[i], sentence[i + 1])
+                self.bigram_counts[bigram] += 1
+            
+            # Count trigrams
+            for i in range(len(sentence) - 2):
+                trigram = (sentence[i], sentence[i + 1], sentence[i + 2])
+                self.trigram_counts[trigram] += 1
+        
+        print(f"Training complete:")
+        print(f"  Vocabulary size: {len(self.vocab):,}")
+        print(f"  Unique unigrams: {len(self.unigram_counts):,}")
+        print(f"  Unique bigrams: {len(self.bigram_counts):,}")
+        print(f"  Unique trigrams: {len(self.trigram_counts):,}")
+    
+    def probability(self, w3: str, w1: str, w2: str) -> float:
+        """
+        Compute P(w3 | w1, w2) with Laplace smoothing.
+        
+        Formula:
+            P(w3|w1,w2) = (C(w1,w2,w3) + α) / (C(w1,w2) + α|V|)
+        
+        Args:
+            w3: The word to predict.
+            w1: First context word (two positions before w3).
+            w2: Second context word (one position before w3).
+            
+        Returns:
+            Conditional probability P(w3 | w1, w2).
+        """
+        trigram_count = self.trigram_counts.get((w1, w2, w3), 0)
+        bigram_count = self.bigram_counts.get((w1, w2), 0)
+        vocab_size = len(self.vocab)
+        
+        # Laplace smoothing
+        numerator = trigram_count + self.smoothing
+        denominator = bigram_count + (self.smoothing * vocab_size)
+        
+        return numerator / denominator if denominator > 0 else 0.0
+    
+    def log_probability(self, w3: str, w1: str, w2: str) -> float:
+        """
+        Compute log P(w3 | w1, w2) for numerical stability.
+        
+        Args:
+            w3: The word to predict.
+            w1: First context word.
+            w2: Second context word.
+            
+        Returns:
+            Log probability.
+        """
+        prob = self.probability(w3, w1, w2)
+        return math.log(prob) if prob > 0 else float('-inf')
+    
+    def predict_next_words(
+        self, 
+        context: str, 
+        top_k: int = 5,
+        exclude_special: bool = True
+    ) -> List[Tuple[str, float]]:
+        """
+        Predict the top-k most likely next words given a context.
+        
+        Args:
+            context: Input text (will use last two words as context).
+            top_k: Number of predictions to return.
+            exclude_special: If True, exclude <s> and </s> from predictions.
+            
+        Returns:
+            List of (word, probability) tuples, sorted by probability descending.
+        """
+        # Tokenize and extract last two words
+        words = context.lower().split()
+        
+        if len(words) == 0:
+            w1, w2 = '<s>', '<s>'
+        elif len(words) == 1:
+            w1, w2 = '<s>', words[0]
+        else:
+            w1, w2 = words[-2], words[-1]
+        
+        # Compute probability for each word in vocabulary
+        candidates = []
+        for word in self.vocab:
+            if exclude_special and word in ('<s>', '</s>'):
+                continue
+            prob = self.probability(word, w1, w2)
+            candidates.append((word, prob))
+        
+        # Sort by probability descending
+        candidates.sort(key=lambda x: x[1], reverse=True)
+        
+        return candidates[:top_k]
+    
+    def sentence_probability(self, tokens: List[str]) -> float:
+        """
+        Compute the probability of a sentence.
+        
+        Args:
+            tokens: Tokenized sentence WITH start/end markers.
+            
+        Returns:
+            Log probability of the sentence.
+        """
+        if len(tokens) < 3:
+            return float('-inf')
+        
+        log_prob = 0.0
+        for i in range(2, len(tokens)):
+            log_prob += self.log_probability(tokens[i], tokens[i-2], tokens[i-1])
+        
+        return log_prob
+    
+    def perplexity(self, sentences: List[List[str]]) -> float:
+        """
+        Compute perplexity on a set of sentences.
+        
+        Perplexity = exp(-1/N * sum(log P(w_i | w_{i-2}, w_{i-1})))
+        
+        Lower perplexity = better model fit.
+        
+        Args:
+            sentences: List of tokenized sentences with boundary markers.
+            
+        Returns:
+            Perplexity score.
+        """
+        total_log_prob = 0.0
+        total_words = 0
+        
+        for sentence in sentences:
+            if len(sentence) < 3:
+                continue
+            for i in range(2, len(sentence)):
+                total_log_prob += self.log_probability(
+                    sentence[i], sentence[i-2], sentence[i-1]
+                )
+                total_words += 1
+        
+        if total_words == 0:
+            return float('inf')
+        
+        avg_log_prob = total_log_prob / total_words
+        return math.exp(-avg_log_prob)
+    
+    def get_context_distribution(
+        self, 
+        w1: str, 
+        w2: str, 
+        top_k: Optional[int] = None
+    ) -> List[Tuple[str, float]]:
+        """
+        Get the probability distribution for a specific bigram context.
+        
+        Args:
+            w1: First context word.
+            w2: Second context word.
+            top_k: If provided, return only top-k predictions.
+            
+        Returns:
+            List of (word, probability) tuples.
+        """
+        candidates = []
+        for word in self.vocab:
+            if word not in ('<s>', '</s>'):
+                prob = self.probability(word, w1, w2)
+                candidates.append((word, prob))
+        
+        candidates.sort(key=lambda x: x[1], reverse=True)
+        
+        if top_k:
+            return candidates[:top_k]
+        return candidates
+    
+    def get_stats(self) -> Dict[str, int]:
+        """
+        Get model statistics.
+        
+        Returns:
+            Dictionary of statistics.
+        """
+        return {
+            'vocab_size': len(self.vocab),
+            'unique_unigrams': len(self.unigram_counts),
+            'unique_bigrams': len(self.bigram_counts),
+            'unique_trigrams': len(self.trigram_counts),
+            'total_tokens': self._total_unigrams,
+        }
+
+
+if __name__ == "__main__":
+    # Quick test with sample data
+    sample_sentences = [
+        ['<s>', '<s>', 'i', 'dey', 'go', 'market', '</s>'],
+        ['<s>', '<s>', 'i', 'dey', 'come', 'back', '</s>'],
+        ['<s>', '<s>', 'you', 'dey', 'go', 'where', '?', '</s>'],
+        ['<s>', '<s>', 'how', 'far', '?', '</s>'],
+        ['<s>', '<s>', 'e', 'don', 'happen', '</s>'],
+    ]
+    
+    model = TrigramLM(smoothing=1.0)
+    model.train(sample_sentences)
+    
+    print("\nTest Predictions:")
+    contexts = ["i dey", "you dey", "how"]
+    for ctx in contexts:
+        preds = model.predict_next_words(ctx, top_k=3)
+        print(f"  '{ctx}' -> {preds}")

src/utils.py

@@ -0,0 +1,85 @@
+"""
+Utility functions for the next-word prediction system.
+"""
+
+from typing import List, Tuple
+import math
+
+
+def format_predictions(predictions: List[Tuple[str, float]], show_percent: bool = True) -> str:
+    """
+    Format prediction results for display.
+    
+    Args:
+        predictions: List of (word, probability) tuples.
+        show_percent: If True, show as percentage.
+        
+    Returns:
+        Formatted string.
+    """
+    lines = []
+    for word, prob in predictions:
+        if show_percent:
+            lines.append(f"  {word}: {prob*100:.2f}%")
+        else:
+            lines.append(f"  {word}: {prob:.6f}")
+    return "\n".join(lines)
+
+
+def calculate_entropy(probabilities: List[float]) -> float:
+    """
+    Calculate entropy of a probability distribution.
+    
+    H(X) = -sum(p * log2(p))
+    
+    Args:
+        probabilities: List of probabilities.
+        
+    Returns:
+        Entropy in bits.
+    """
+    entropy = 0.0
+    for p in probabilities:
+        if p > 0:
+            entropy -= p * math.log2(p)
+    return entropy
+
+
+def top_k_accuracy(
+    model, 
+    test_sentences: List[List[str]], 
+    k: int = 5
+) -> float:
+    """
+    Calculate top-k accuracy on test data.
+    
+    Measures what fraction of true next words appear in top-k predictions.
+    
+    Args:
+        model: Trained TrigramLM instance.
+        test_sentences: List of tokenized sentences with markers.
+        k: Number of top predictions to consider.
+        
+    Returns:
+        Accuracy as fraction between 0 and 1.
+    """
+    correct = 0
+    total = 0
+    
+    for sentence in test_sentences:
+        if len(sentence) < 3:
+            continue
+        
+        for i in range(2, len(sentence)):
+            w1, w2 = sentence[i-2], sentence[i-1]
+            true_word = sentence[i]
+            
+            # Get top-k predictions
+            preds = model.get_context_distribution(w1, w2, top_k=k)
+            pred_words = [w for w, _ in preds]
+            
+            if true_word in pred_words:
+                correct += 1
+            total += 1
+    
+    return correct / total if total > 0 else 0.0