jeevan0704/sequential-model-for-sequential_dataset

1import torch2import numpy as np3import os4import config5from model import create_model6from utils import pad_sequences, create_padding_mask7 8class ModelInference:9    def __init__(self):10        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')11        self.model = None12        self.vocab_size = 32737  # Fixed from training analysis13        self.max_length = config.MAX_SEQ_LENGTH14        self.model_path = os.path.join(config.MODEL_DIR, 'best_model.pth')15        self._load_model()16        17    def _load_model(self):18        """Load the trained model"""19        try:20            print(f"Loading model architecture with vocab_size={self.vocab_size}...")21            self.model = create_model(self.vocab_size)22            23            if os.path.exists(self.model_path):24                print(f"Loading weights from {self.model_path}...")25                self.model.load_state_dict(torch.load(self.model_path, map_location=self.device))26                print("Model loaded successfully!")27            else:28                print(f"WARNING: Model file not found at {self.model_path}")29                print("Initializing with random weights for TESTING PURPOSES ONLY.")30                31            self.model.to(self.device)32            self.model.eval()33            34        except Exception as e:35            print(f"Error loading model: {str(e)}")36            raise e37 38    def predict(self, sequence_str):39        """40        Run inference on a comma-separated string of integers41        42        Args:43            sequence_str: String like "1, 45, 23"44            45        Returns:46            List of predicted labels for the valid part of the sequence47        """48        if not sequence_str or not sequence_str.strip():49            return []50            51        try:52            # Parse input string to integers53            # integers = [int(x.strip()) for x in sequence_str.split(',') if x.strip()]54            55            # Handle user input more robustly - replace non-numeric chars with space?56            # Or just split by comma/space57            import re58            tokens = re.findall(r'\d+', sequence_str)59            input_seq = [int(t) for t in tokens]60            61            if not input_seq:62                return []63                64            original_length = len(input_seq)65            66            # Truncate if too long (or split? let's truncate for now)67            if len(input_seq) > self.max_length:68                input_seq = input_seq[:self.max_length]69                70            # Pad (returns 2D array [batch_size, max_len])71            padded_batch = pad_sequences([input_seq], self.max_length, padding_value=0)72            73            # Create mask (returns 2D array [batch_size, max_len])74            mask_batch = create_padding_mask([input_seq], self.max_length)75            76            # Convert to tensor77            input_tensor = torch.tensor(padded_batch, dtype=torch.long).to(self.device)78            mask_tensor = torch.tensor(mask_batch, dtype=torch.bool).to(self.device)79            80            # Inference81            with torch.no_grad():82                logits = self.model(input_tensor, mask_tensor)83                predictions = torch.argmax(logits, dim=2)84                85            # Get predictions for valid length (unpadded)86            # Prediction shape: [1, seq_len, num_classes] -> argmax -> [1, seq_len]87            preds = predictions[0].cpu().numpy()88            89            # Remap labels {0, 1, 2} back to {-1, 0, 1}90            # 0 -> -1, 1 -> 0, 2 -> 191            final_predictions = preds - 192            93            # Return only the valid part corresponding to input length94            # Note: input_seq might have been truncated95            valid_len = min(original_length, self.max_length)96            return final_predictions[:valid_len].tolist()97            98        except Exception as e:99            print(f"Inference error: {str(e)}")100            return []101 102if __name__ == "__main__":103    # Test inference104    infer = ModelInference()105    sample = "1, 10, 50, 100"106    preds = infer.predict(sample)107    print(f"Input: {sample}")108    print(f"Predictions: {preds}")109