pragunk/PropagationShield

1import torch2from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig3from typing import Dict, List, Any4 5class EndpointHandler:6    def __init__(self, path=""):7        """8        Initializes the model and tokenizer. 9        `path` is automatically provided by Hugging Face (it points to your repo files).10        """11        print("🚀 Initializing PropagationShield Handler...")12        13        self.tokenizer = AutoTokenizer.from_pretrained(path)14        15        # 1. Configure 4-bit quantization to prevent OOM and System RAM limits16        bnb_config = BitsAndBytesConfig(17            load_in_4bit=True,18            bnb_4bit_use_double_quant=True,19            bnb_4bit_quant_type="nf4",20            bnb_4bit_compute_dtype=torch.float1621        )22        23        # 2. Load the model safely24        self.model = AutoModelForCausalLM.from_pretrained(25            path,26            quantization_config=bnb_config,27            device_map="auto",28            torch_dtype=torch.float16,29            low_cpu_mem_usage=True, # Crucial to prevent the 30GB RAM crash during boot30        )31        print("✅ PropagationShield Loaded Successfully!")32 33    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:34        """35        Runs inference on the incoming request.36        """37        # Parse incoming data38        inputs = data.pop("inputs", data)39        parameters = data.pop("parameters", {})40        41        max_new_tokens = parameters.get("max_new_tokens", 512)42        temperature = parameters.get("temperature", 0.1)43        44        # 3. Format the prompt45        # If the user sends a list of messages [{"role": "system", "content": "..."}, ...]46        if isinstance(inputs, list):47            prompt = self.tokenizer.apply_chat_template(48                inputs, tokenize=False, add_generation_prompt=True49            )50        # If the user sends a raw formatted string51        else:52            prompt = str(inputs)53            54        # 4. Tokenize55        input_ids = self.tokenizer(prompt, return_tensors="pt").input_ids.to(self.model.device)56        57        # 5. Generate58        with torch.no_grad():59            output_ids = self.model.generate(60                input_ids,61                max_new_tokens=max_new_tokens,62                temperature=temperature,63                do_sample=True if temperature > 0.0 else False,64                pad_token_id=self.tokenizer.eos_token_id65            )66            67        # 6. Isolate and decode only the newly generated tokens68        generated_ids = output_ids[0][input_ids.shape[-1]:]69        generated_text = self.tokenizer.decode(generated_ids, skip_special_tokens=True)70        71        # Return in standard HF API format72        return [{"generated_text": generated_text.strip()}]