def predict_proba()

in src/infer_location.py [0:0]

• 12 lines of code
• 1 McCabe index (conditional complexity)

    def predict_proba(self, texts):
        # Set max_length to the model's maximum output sequence length
        max_length = 64
        num_classes = self.ort_session.get_outputs()[0].shape[-1]  # Number of classes from the model

    
        # Initialize 3D array for all probabilities with zeros (padding)
        all_probs = np.zeros((len(texts), max_length, num_classes))
        
        for i, text in enumerate(texts):
            # Tokenize and prepare input
            input_ids, attention_mask = self.preprocess(text)
            
            # Run model inference to get logits
            logits = self.predict((input_ids, attention_mask))
            
            # Apply softmax to get probabilities for each token
            prob = F.softmax(torch.tensor(logits), dim=-1).numpy()
            
            # Trim to actual token length, ignoring padding tokens
            tokens = self.tokenizer(text, truncation=True, padding="max_length")['input_ids']
            token_probabilities = prob[:len(tokens), :]
            
            # Copy token probabilities into the pre-initialized array up to token length
            all_probs[i, :len(tokens), :] = token_probabilities  # Fill up to token length
        
        # Return the padded 3D numpy array
        return all_probs