joseAndres777/WazapSplitter-LLM

1from transformers import AutoTokenizer, AutoModelForCausalLM2from peft import PeftModel3import torch4import json5import os6 7os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"8 9class EndpointHandler:10    def __init__(self, path=""):11        """12        Initialize the handler with the model from the given path13        """14        model_name = "meta-llama/Llama-3.3-70B-Instruct"15 16        self.tokenizer = AutoTokenizer.from_pretrained(model_name)17 18        base_model = AutoModelForCausalLM.from_pretrained(19            model_name,20            torch_dtype=torch.float16,21            device_map="auto",22            trust_remote_code=True,23            load_in_8bit=True,  24            low_cpu_mem_usage=True25        )26 27        try:28            self.model = PeftModel.from_pretrained(29                base_model,30                path,31                is_trainable=False 32            )33            print("Successfully loaded adapter with base model")34        except Exception as e:35            print(f"Error loading adapter: {e}")36            print("Falling back to base model without adapter")37            self.model = base_model38        39        try:40            with open(f"{path}/chat_template.jinja", "r") as f:41                self.chat_template = f.read()42        except:43            self.chat_template = None44 45    def __call__(self, data):46        """47        Process the input data and return the model's response48        """49        inputs = data.get("inputs", "")50        parameters = data.get("parameters", {})51        52        default_prompt = "Break this text into WhatsApp messages like a real person would send them. Split where you'd naturally pause: after greetings, before/after questions, between different thoughts, when changing topics. Preserve exact wording - just divide where someone would actually hit 'send' and start a new message. Output JSON array."53 54        custom_prompt = parameters.get("prompt", default_prompt)55 56        messages = [57            {"role": "system", "content": custom_prompt},58            {"role": "user", "content": inputs}59        ]60        61        if self.chat_template:62            text = self.tokenizer.apply_chat_template(63                messages, 64                tokenize=False, 65                add_generation_prompt=True66            )67        else:68            text = f"{custom_prompt}\nUser: {inputs}\nAssistant:"69        70        # Tokenize71        model_inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)72        73        # Generate response74        with torch.no_grad():75            outputs = self.model.generate(76                **model_inputs,77                max_new_tokens=parameters.get("max_new_tokens", 100),78                temperature=parameters.get("temperature", 0.3),  79                top_p=parameters.get("top_p", 0.9),80                do_sample=True,81                pad_token_id=self.tokenizer.eos_token_id,82                repetition_penalty=1.183            )84        85    86        response = self.tokenizer.decode(87            outputs[0][model_inputs.input_ids.shape[-1]:],88            skip_special_tokens=True89        ).strip()90 91        try:92          93            if response.startswith('[') and response.endswith(']'):94                parsed = json.loads(response)95                if isinstance(parsed, list):96                    formatted_response = response97                else:98                    formatted_response = json.dumps([response])99            else:100 101                formatted_response = json.dumps([response])102        except:103            formatted_response = json.dumps([inputs])104 105        return [{"content": formatted_response}]