File size: 2,494 Bytes
70b9491 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
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, Tira Beauty, netmeds and milkbasket. 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):] |