llama_3_1_8B_4bit / handler.py
sundeepsingh01's picture
Update handler.py
76e0eed verified
Raw
History Blame Contribute Delete
1.42 kB
from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace, HuggingFaceEndpoint
import torch
from huggingface_hub import login
api_key = 'hf'+ '_' + 'tlVzheuQBwjAxOtNKPqnHSQprFYnDLllut'
login(token=api_key)
class EndpointHandler:
def __init__(self, path1="Cognute02/llama_3_1_8B_4bit"):
# load model and processor from path
quantization_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype="float16",
bnb_4bit_use_double_quant=True,
)
llm = HuggingFacePipeline.from_model_id(
model_id=path1,
task="text-generation",
pipeline_kwargs=dict(
max_new_tokens=512,
do_sample=False,
repetition_penalty=1.03,
return_full_text=False,
temperature = 0.25
),
model_kwargs={"quantization_config": quantization_config},
)
self.chatllm = ChatHuggingFace(llm=llm)
def __call__(self, data):
inputs = data['inputs']
tools = data['tools']
llm_ = self.chatllm.bind_tools(tools)
outputs = llm_.invoke(inputs)
return outputs