wins0nLYQ/fraud-message-detector

1import torch2 3from peft import PeftModel4from typing import List, Any, Dict5from transformers import AutoModelForSequenceClassification, AutoTokenizer, BitsAndBytesConfig6 7class EndpointHandler():8    def __init__(self, path=""):9        # Load the model and tokenizer10        self.tokenizer = AutoTokenizer.from_pretrained(path)11 12        quantization_config = BitsAndBytesConfig(13            load_in_4bit=True,14            bnb_4bit_use_double_quant=True,15            bnb_4bit_quant_type="nf4",16            bnb_4bit_compute_dtype=torch.bfloat1617        )18        19        self.model = AutoModelForSequenceClassification.from_pretrained(path, quantization_config=quantization_config, num_labels=2, device_map="auto")20        self.peft_model = PeftModel.from_pretrained(self.model, path)21        22    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:23        """24        Args:25            inputs (:obj: 'str')26        Return:27            A :obj: 'list' | 'dict'28        """29        30        # Extract the input text from the data31        input_text = data.pop("inputs", data)32        33        if not input_text:34            return [{"error": "No input text provided."}]35 36        # Tokenize input text37        encoded_input = self.tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=512)38 39        # Perform inference40        with torch.no_grad():41            outputs = self.peft_model(**encoded_input)42 43        # Process output (assuming this is a classification model)44        scores = torch.nn.functional.softmax(outputs.logits, dim=-1)45        predicted_class = torch.argmax(scores, dim=-1).item()46 47        # Return the prediction and scores as an array48        return [{"predicted_class": predicted_class}]49