Create handler.py

handler.py

@@ -0,0 +1,69 @@
+import time
+import bitsandbytes as bnb
+import torch
+import transformers
+from datasets import load_dataset
+from typing import Dict, List, Any
+from peft import (
+    LoraConfig,
+    PeftConfig,
+    PeftModel,
+    get_peft_model,
+    prepare_model_for_kbit_training,
+)
+from transformers import (
+    AutoConfig,
+    LlamaTokenizer,
+    LlamaForCausalLM,
+    #AutoModelForCausalLM,
+    #AutoTokenizer,
+    BitsAndBytesConfig,
+)
+import json
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+)
+
+
+from huggingface_hub import login
+access_token_read = "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"
+login(token = access_token_read)
+
+
+class EndpointHandler:
+    def __init__(self, path=''):
+        PEFT_MODEL = path
+        config = PeftConfig.from_pretrained(PEFT_MODEL)
+        self.model = LlamaForCausalLM.from_pretrained(
+            config.base_model_name_or_path,
+            return_dict=True,
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+        )
+        self.tokenizer = LlamaTokenizer.from_pretrained(config.base_model_name_or_path)
+        self.tokenizer.pad_token_id = (0)
+        self.tokenizer.padding_side = "left"
+        self.model = PeftModel.from_pretrained(self.model, PEFT_MODEL)
+        self.generation_config = self.model.generation_config
+        self.generation_config.max_new_tokens = 100
+        self.generation_config.pad_token_id = self.tokenizer.eos_token_id
+        self.generation_config.eos_token_id = self.tokenizer.eos_token_id
+
+
+    def __call__(self, data: Dict[str, Any]):
+        prompt = data.pop("inputs", data)
+        DEVICE = "cuda:0"
+        input_message = f"""[INST]You are an assistant that detects the intent and entity of user's message. Possible entity stores are JioMart, JioFiber, JioCinema and Tira Beauty. Detect the intent and entity of the following user's message[/INST]\nUser: {prompt}\nAssistant: """.strip()
+        encoding = self.tokenizer(input_message, return_tensors="pt").to(DEVICE)
+        with torch.inference_mode():
+            outputs = self.model.generate(
+                input_ids=encoding.input_ids,
+                attention_mask=encoding.attention_mask,
+                generation_config=self.generation_config
+            )
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)[len(input_message):]