llama_3_1_8B_4bit / handler_old.py
sundeepsingh01's picture
Rename handler.py to handler_old.py
71eb73d verified
Raw
History Blame Contribute Delete
1.4 kB
from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace, HuggingFaceEndpoint
import torch
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,
)
self.model = AutoModelForCausalLM.from_pretrained(path1, quantization_config=quantization_config)
self.tokenizer = AutoTokenizer.from_pretrained(path1)
self.pipe = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer)
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Args:
data (:obj:):
includes the deserialized image file as PIL.Image
"""
inputs = data['inputs']
outputs = self.pipe(inputs, max_new_tokens=256, do_sample=False, return_full_text=False,
temperature=0.25, top_p=0.25, top_k=10)
return outputs