Add TinyBert.py

TinyBert.py

@@ -0,0 +1,620 @@
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+from torch.utils.data import Dataset
+import re
+
+
+class IntentDataset(Dataset):
+    """
+    Dataset for handling student input and session context for 5-class intent categorization.
+    """
+    def __init__(self, data, tokenizer, max_length=128):
+        # data: list of dicts with 'student_input', 'session_context', 'label'
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        self.label_map = {
+            'On-Topic Question': 0,
+            'Off-Topic Question': 1,
+            'Emotional-State': 2,
+            'Pace-Related': 3,
+            'Repeat/clarification': 4
+        }
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        item = self.data[idx]
+        student_input = str(item.get('student_input', ''))
+        session_context = str(item.get('session_context', ''))
+        
+        # Tokenize pair — longest_first truncation preserves student input priority
+        encoded = self.tokenizer(
+            student_input,
+            session_context,
+            padding='max_length',
+            truncation='longest_first',
+            max_length=self.max_length,
+            return_tensors='pt'
+        )
+        
+        label_val = item.get('label', 0)
+        if isinstance(label_val, str):
+            label_val = self.label_map.get(label_val, 0)
+            
+        output = {
+            'input_ids': encoded['input_ids'].squeeze(0),
+            'attention_mask': encoded['attention_mask'].squeeze(0),
+            'labels': torch.tensor(label_val, dtype=torch.long)
+        }
+        if 'token_type_ids' in encoded:
+            output['token_type_ids'] = encoded['token_type_ids'].squeeze(0)
+            
+        return output
+
+
+class CompoundSentenceSplitter:
+    """
+    Algorithm to split compound sentences containing 2 separate questions.
+    Handles various patterns and conjunctions commonly used to combine questions.
+    English only.
+    """
+    
+    def __init__(self):
+        # English question words
+        self.question_words = [
+            'what', 'when', 'where', 'which', 'who', 'whom', 'whose', 'why', 'how',
+            'is', 'are', 'was', 'were', 'do', 'does', 'did', 'can', 'could', 
+            'will', 'would', 'should', 'may', 'might', 'must'
+        ]
+        
+        # English conjunctions
+        self.conjunctions = [
+            'and', 'or', 'also', 'plus', 'additionally', 'moreover'
+        ]
+        
+        # English transition phrases
+        self.transition_phrases = [
+            'and also', 'and what about', 'and how about', 'or what about', 
+            'or how about', 'also what', 'also how', 'also when', 'also where',
+            'also who', 'also why', 'plus what', 'plus how'
+        ]
+    
+    def split_compound_question(self, text):
+        """
+        Split a compound sentence into 2 separate questions if applicable.
+        Works with English text.
+        
+        Args:
+            text (str): Input text that may contain compound questions
+            
+        Returns:
+            list: List of separated questions. Returns [text] if no split is needed.
+        """
+        text = text.strip()
+        
+        # Check if text is likely a question
+        if not self._is_question(text):
+            return [text]
+        
+        # Try different splitting strategies
+        questions = []
+        
+        # Strategy 1: Split by transition phrases
+        questions = self._split_by_transition_phrases(text)
+        if len(questions) > 1:
+            return self._clean_questions(questions)
+        
+        # Strategy 2: Split by conjunction followed by question word
+        questions = self._split_by_conjunction_pattern(text)
+        if len(questions) > 1:
+            return self._clean_questions(questions)
+        
+        # Strategy 3: Split by semicolon or comma-conjunction pattern
+        questions = self._split_by_punctuation_pattern(text)
+        if len(questions) > 1:
+            return self._clean_questions(questions)
+        
+        # Strategy 4: Split by multiple question marks
+        questions = self._split_by_question_marks(text)
+        if len(questions) > 1:
+            return self._clean_questions(questions)
+        
+        # No split found, return original
+        return [text]
+    
+    def _is_question(self, text):
+        """Check if text is likely a question (English)"""
+        text_stripped = text.strip()
+        
+        # Has question mark
+        if '?' in text:
+            return True
+        
+        # Check for question words at the start
+        words = text_stripped.split()
+        if words:
+            first_word = words[0].lower()
+            # Check English question words
+            if first_word in self.question_words:
+                return True
+        
+        return False
+    
+    def _split_by_transition_phrases(self, text):
+        """Split by transition phrases (English)"""
+        for phrase in self.transition_phrases:
+            # English phrase with word boundaries
+            pattern = r'\s+' + re.escape(phrase) + r'\s+'
+            
+            match = re.search(pattern, text, re.IGNORECASE)
+            if match:
+                parts = re.split(pattern, text, maxsplit=1, flags=re.IGNORECASE)
+                if len(parts) == 2 and parts[0] and parts[1]:
+                    return parts
+        
+        return [text]
+    
+    def _split_by_conjunction_pattern(self, text):
+        """Split by conjunction followed by question word (English)"""
+        # Pattern: conjunction + question word
+        for conj in self.conjunctions:
+            for qword in self.question_words:
+                # English pattern with word boundaries
+                pattern = r'\s+' + re.escape(conj) + r'\s+' + re.escape(qword) + r'\b'
+                
+                match = re.search(pattern, text, re.IGNORECASE)
+                
+                if match:
+                    # Find the actual position in original text
+                    split_pos = match.start()
+                    part1 = text[:split_pos].strip()
+                    part2 = text[split_pos:].strip()
+                    
+                    # Remove leading conjunction from part2
+                    for c in self.conjunctions:
+                        is_arabic_c = any(ch in 'أبتثجحخدذرزسشصضطظعغفقكلمنهويىةؤإآ' for ch in c)
+                        part2 = re.sub(r'^\s*' + re.escape(c) + r'\s+', '', part2, flags=re.IGNORECASE if not is_arabic_c else 0)
+                    
+                    # Ensure both parts are questions
+                    if part1 and part2 and self._is_question(part1):
+                        return [part1, part2]
+        
+        return [text]
+    
+    def _split_by_punctuation_pattern(self, text):
+        """Split by semicolon or specific comma patterns"""
+        # Split by semicolon (works for both languages)
+        if ';' in text or '؛' in text:  # Added Arabic semicolon
+            parts = re.split(r'[;؛]', text, maxsplit=1)
+            if len(parts) == 2:
+                parts = [p.strip() for p in parts]
+                if all(self._is_question(p) for p in parts):
+                    return parts
+        
+        # Split by comma followed by question word
+        pattern = r',\s+(?=' + '|'.join([re.escape(qw) for qw in self.question_words]) + r')'
+        parts = re.split(pattern, text, maxsplit=1, flags=re.IGNORECASE)
+        
+        if len(parts) == 2:
+            parts = [p.strip() for p in parts]
+            # Only split if second part is clearly a question
+            if self._is_question(parts[1]):
+                return parts
+        
+        return [text]
+    
+    def _split_by_question_marks(self, text):
+        """Split by question marks if multiple exist (both ? and ؟)"""
+        # Count both English and Arabic question marks
+        q_marks = text.count('?') + text.count('؟')
+        
+        if q_marks >= 2:
+            # Split at first question mark
+            match = re.search(r'[?؟]', text)
+            if match:
+                split_pos = match.end()
+                part1 = text[:split_pos].strip()
+                part2 = text[split_pos:].strip()
+                
+                if part2:  # Ensure second part is not empty
+                    return [part1, part2]
+        
+        return [text]
+    
+    def _clean_questions(self, questions):
+        """Clean and validate split questions"""
+        cleaned = []
+        
+        for q in questions:
+            q = q.strip()
+            
+            # Skip empty questions
+            if not q:
+                continue
+            
+            # Ensure question ends with '?' or '؟' if it's clearly a question
+            if self._is_question(q):
+                # Check if already has question mark
+                if not (q.endswith('?') or q.endswith('؟')):
+                    # Add appropriate question mark based on language
+                    if any(c in 'أبتثجحخدذرزسشصضطظعغفقكلمنهويىةؤإآ' for c in q):
+                        q += '؟'  # Arabic question mark
+                    else:
+                        q += '?'  # English question mark
+            
+            cleaned.append(q)
+        
+        return cleaned if len(cleaned) > 1 else [' '.join(questions)]
+
+
+class TinyBertCNN(nn.Module):
+    """
+    TinyBERT-CNN model for intent classification.
+    Combines TinyBERT embeddings with CNN layers + BatchNorm + hidden FC layer.
+    """
+    
+    def __init__(
+        self,
+        num_classes,
+        bert_model_name='huawei-noah/TinyBERT_General_4L_312D',
+        num_filters=256,
+        filter_sizes=[2, 3, 4],
+        dropout=0.5,
+        hidden_dim=128,
+        freeze_bert=False
+    ):
+        """
+        Args:
+            num_classes (int): Number of intent classes
+            bert_model_name (str): Pre-trained TinyBERT model name
+            num_filters (int): Number of filters for each filter size
+            filter_sizes (list): List of filter sizes for CNN
+            dropout (float): Dropout rate
+            hidden_dim (int): Hidden FC layer dimension
+            freeze_bert (bool): Whether to freeze BERT parameters
+        """
+        super(TinyBertCNN, self).__init__()
+        
+        # Load TinyBERT model
+        self.bert = AutoModel.from_pretrained(bert_model_name)
+        self.bert_hidden_size = self.bert.config.hidden_size
+        
+        # Freeze BERT parameters if specified
+        if freeze_bert:
+            for param in self.bert.parameters():
+                param.requires_grad = False
+        
+        # CNN layers with BatchNorm
+        self.convs = nn.ModuleList([
+            nn.Conv1d(
+                in_channels=self.bert_hidden_size,
+                out_channels=num_filters,
+                kernel_size=fs
+            )
+            for fs in filter_sizes
+        ])
+        self.batchnorms = nn.ModuleList([
+            nn.BatchNorm1d(num_filters)
+            for _ in filter_sizes
+        ])
+        
+        # Dropout
+        self.dropout = nn.Dropout(dropout)
+        
+        # Hidden FC layer
+        cnn_out_dim = len(filter_sizes) * num_filters
+        self.fc_hidden = nn.Linear(cnn_out_dim, hidden_dim)
+        self.bn_hidden = nn.BatchNorm1d(hidden_dim)
+        
+        # Output layer
+        self.fc = nn.Linear(hidden_dim, num_classes)
+        
+    def forward(self, input_ids, attention_mask, token_type_ids=None):
+        """
+        Forward pass
+        
+        Args:
+            input_ids: Token IDs (batch_size, seq_len)
+            attention_mask: Attention mask (batch_size, seq_len)
+            token_type_ids: Token type IDs (batch_size, seq_len), optional
+            
+        Returns:
+            logits: Classification logits (batch_size, num_classes)
+        """
+        # Get TinyBERT embeddings
+        # outputs: (batch_size, seq_len, hidden_size)
+        bert_kwargs = {
+            'input_ids': input_ids,
+            'attention_mask': attention_mask
+        }
+        if token_type_ids is not None:
+            bert_kwargs['token_type_ids'] = token_type_ids
+            
+        bert_output = self.bert(**bert_kwargs)
+        
+        # Use last hidden state
+        # sequence_output: (batch_size, seq_len, hidden_size)
+        sequence_output = bert_output.last_hidden_state
+        
+        # Transpose for CNN: (batch_size, hidden_size, seq_len)
+        sequence_output = sequence_output.transpose(1, 2)
+        
+        # Pad if sequence is shorter than the largest kernel
+        max_kernel = max(conv.kernel_size[0] for conv in self.convs)
+        if sequence_output.size(2) < max_kernel:
+            pad_size = max_kernel - sequence_output.size(2)
+            sequence_output = torch.nn.functional.pad(sequence_output, (0, pad_size))
+        
+        # Apply convolution + batchnorm + max pooling for each filter size
+        conv_outputs = []
+        for conv, bn in zip(self.convs, self.batchnorms):
+            # conv_out: (batch_size, num_filters, seq_len - filter_size + 1)
+            conv_out = torch.relu(bn(conv(sequence_output)))
+            # pooled: (batch_size, num_filters)
+            pooled = torch.max_pool1d(conv_out, conv_out.size(2)).squeeze(2)
+            conv_outputs.append(pooled)
+        
+        # Concatenate all features
+        # concatenated: (batch_size, len(filter_sizes) * num_filters)
+        concatenated = torch.cat(conv_outputs, dim=1)
+        concatenated = self.dropout(concatenated)
+        
+        # Hidden FC layer
+        hidden = torch.relu(self.bn_hidden(self.fc_hidden(concatenated)))
+        hidden = self.dropout(hidden)
+        
+        # Final classification
+        logits = self.fc(hidden)
+        
+        return logits
+
+
+class IntentClassifier:
+    """
+    Wrapper class for training and inference
+    """
+    
+    def __init__(
+        self,
+        num_classes,
+        bert_model_name='huawei-noah/TinyBERT_General_4L_312D',
+        num_filters=256,
+        filter_sizes=[2, 3, 4],
+        dropout=0.5,
+        freeze_bert=False,
+        device=None
+    ):
+        self.device = device or torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        
+        # Initialize model
+        self.model = TinyBertCNN(
+            num_classes=num_classes,
+            bert_model_name=bert_model_name,
+            num_filters=num_filters,
+            filter_sizes=filter_sizes,
+            dropout=dropout,
+            freeze_bert=freeze_bert
+        ).to(self.device)
+        
+        # Initialize tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
+        
+        # Initialize compound sentence splitter
+        self.sentence_splitter = CompoundSentenceSplitter()
+        
+        self.num_classes = num_classes
+        
+    def preprocess_text(self, text):
+        """
+        Preprocess text by splitting compound questions if detected
+        
+        Args:
+            text (str): Input text (English or Arabic)
+            
+        Returns:
+            list: List of individual questions
+        """
+        return self.sentence_splitter.split_compound_question(text)
+    
+    def predict(self, student_inputs, session_contexts=None, max_length=128, split_compound=False):
+        """
+        Predict intents for input texts
+        
+        Args:
+            student_inputs (list): List of student input texts (English or Arabic)
+            session_contexts (list): List of session context texts
+            max_length (int): Maximum sequence length
+            split_compound (bool): Whether to split compound questions before prediction
+            
+        Returns:
+            If split_compound=False:
+                predictions: Predicted class indices
+                probabilities: Prediction probabilities
+            If split_compound=True:
+                predictions: List of predictions (may contain multiple per text if split)
+                probabilities: List of probabilities
+                split_info: Dictionary with information about splits
+        """
+        # Handle compound questions if requested
+        if split_compound:
+            return self._predict_with_splitting(student_inputs, session_contexts, max_length)
+        
+        self.model.eval()
+        
+        # Determine if we are passing single string or pair
+        if session_contexts is not None:
+            text_args = (student_inputs, session_contexts)
+        else:
+            text_args = (student_inputs,)
+        
+        # Tokenize
+        encoded = self.tokenizer(
+            *text_args,
+            padding=True,
+            truncation=True,
+            max_length=max_length,
+            return_tensors='pt'
+        )
+        
+        input_ids = encoded['input_ids'].to(self.device)
+        attention_mask = encoded['attention_mask'].to(self.device)
+        token_type_ids = encoded.get('token_type_ids')
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.to(self.device)
+        
+        with torch.no_grad():
+            logits = self.model(input_ids, attention_mask, token_type_ids=token_type_ids)
+            probabilities = torch.softmax(logits, dim=1)
+            predictions = torch.argmax(probabilities, dim=1)
+        
+        return predictions.cpu().numpy(), probabilities.cpu().numpy()
+    
+    def _predict_with_splitting(self, student_inputs, session_contexts=None, max_length=128):
+        """
+        Predict intents after splitting compound questions (English and Arabic)
+        
+        Args:
+            student_inputs (list): List of input texts
+            session_contexts (list): List of session context texts
+            max_length (int): Maximum sequence length
+            
+        Returns:
+            predictions: List of predictions (one per original text, may contain multiple if split)
+            probabilities: List of probabilities
+            split_info: Dictionary with information about splits
+        """
+        all_predictions = []
+        all_probabilities = []
+        split_info = {
+            'original_texts': student_inputs,
+            'split_texts': [],
+            'was_split': [],
+            'split_indices': []  # Maps split question index to original text index
+        }
+        
+        # Collect all questions after splitting
+        all_questions = []
+        all_contexts = []
+        for i, text in enumerate(student_inputs):
+            questions = self.preprocess_text(text)
+            split_info['split_texts'].append(questions)
+            split_info['was_split'].append(len(questions) > 1)
+            
+            # Track which original text each split question belongs to
+            for _ in questions:
+                split_info['split_indices'].append(i)
+                if session_contexts is not None:
+                    all_contexts.append(session_contexts[i])
+            
+            all_questions.extend(questions)
+        
+        # Predict for all questions at once
+        if all_questions:
+            contexts_to_pass = all_contexts if session_contexts is not None else None
+            predictions, probabilities = self.predict(all_questions, contexts_to_pass, max_length, split_compound=False)
+            
+            # Reorganize results by original text
+            idx = 0
+            for i, text in enumerate(student_inputs):
+                num_questions = len(split_info['split_texts'][i])
+                text_predictions = predictions[idx:idx + num_questions]
+                text_probabilities = probabilities[idx:idx + num_questions]
+                
+                all_predictions.append(text_predictions)
+                all_probabilities.append(text_probabilities)
+                
+                idx += num_questions
+        
+        return all_predictions, all_probabilities, split_info
+    
+    def train_step(self, batch, optimizer, criterion):
+        """
+        Single training step
+        
+        Args:
+            batch: Dictionary with 'input_ids', 'attention_mask', 'labels'
+            optimizer: Optimizer
+            criterion: Loss function
+            
+        Returns:
+            loss: Training loss
+        """
+        self.model.train()
+        
+        input_ids = batch['input_ids'].to(self.device)
+        attention_mask = batch['attention_mask'].to(self.device)
+        labels = batch['labels'].to(self.device)
+        token_type_ids = batch.get('token_type_ids')
+        if token_type_ids is not None:
+            token_type_ids = token_type_ids.to(self.device)
+        
+        # Forward pass
+        logits = self.model(input_ids, attention_mask, token_type_ids=token_type_ids)
+        loss = criterion(logits, labels)
+        
+        # Backward pass
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        
+        return loss.item()
+    
+    def evaluate(self, dataloader, criterion):
+        """
+        Evaluate model on validation/test set
+        
+        Args:
+            dataloader: DataLoader for evaluation
+            criterion: Loss function
+            
+        Returns:
+            avg_loss: Average loss
+            accuracy: Classification accuracy
+        """
+        self.model.eval()
+        
+        total_loss = 0
+        total_correct = 0
+        total_samples = 0
+        
+        with torch.no_grad():
+            for batch in dataloader:
+                input_ids = batch['input_ids'].to(self.device)
+                attention_mask = batch['attention_mask'].to(self.device)
+                labels = batch['labels'].to(self.device)
+                token_type_ids = batch.get('token_type_ids')
+                if token_type_ids is not None:
+                    token_type_ids = token_type_ids.to(self.device)
+                
+                # Forward pass
+                logits = self.model(input_ids, attention_mask, token_type_ids=token_type_ids)
+                loss = criterion(logits, labels)
+                
+                # Calculate metrics
+                predictions = torch.argmax(logits, dim=1)
+                total_loss += loss.item() * labels.size(0)
+                total_correct += (predictions == labels).sum().item()
+                total_samples += labels.size(0)
+        
+        avg_loss = total_loss / total_samples
+        accuracy = total_correct / total_samples
+        
+        return avg_loss, accuracy
+    
+    def save_model(self, path):
+        """Save model checkpoint"""
+        torch.save({
+            'model_state_dict': self.model.state_dict(),
+            'num_classes': self.num_classes
+        }, path)
+        print(f"Model saved to {path}")
+    
+    def load_model(self, path):
+        """Load model checkpoint"""
+        checkpoint = torch.load(path, map_location=self.device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        print(f"Model loaded from {path}")
+