|
from typing import Dict, List, Any |
|
import transformers |
|
import torch |
|
from datetime import datetime |
|
from transformers import StoppingCriteria, StoppingCriteriaList |
|
from transformers.utils import logging |
|
|
|
logging.set_verbosity_info() |
|
logger = logging.get_logger("transformers") |
|
|
|
|
|
class EndpointHandler(): |
|
|
|
def __init__(self, path=""): |
|
logger.info("111111111111111111111111111") |
|
logger.info(f"Hugging face handler path {path}") |
|
path = 'mosaicml/mpt-7b-instruct' |
|
|
|
self.model = transformers.AutoModelForCausalLM.from_pretrained(path, |
|
|
|
|
|
|
|
trust_remote_code=True, |
|
torch_dtype=torch.bfloat16, |
|
max_seq_len=32000 |
|
) |
|
|
|
self.tokenizer = transformers.AutoTokenizer.from_pretrained('EleutherAI/gpt-neox-20b') |
|
print("tokenizer created ", datetime.now()) |
|
|
|
|
|
stop_token_ids = self.tokenizer.convert_tokens_to_ids(["<|endoftext|>"]) |
|
|
|
class StopOnTokens(StoppingCriteria): |
|
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs): |
|
for stop_id in stop_token_ids: |
|
if input_ids[0][-1] == stop_id: |
|
return True |
|
return False |
|
|
|
stopping_criteria = StoppingCriteriaList([StopOnTokens()]) |
|
|
|
self.generate_text = transformers.pipeline( |
|
model=self.model, |
|
tokenizer=self.tokenizer, |
|
stopping_criteria=stopping_criteria, |
|
task='text-generation', |
|
return_full_text=True, |
|
temperature=0.1, |
|
top_p=0.15, |
|
top_k=0, |
|
max_new_tokens=2048, |
|
repetition_penalty=1.1 |
|
) |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
logger.info(f"iiinnnnnnnnnn {data}") |
|
inputs = data.pop("inputs",data) |
|
logger.info(f"iiinnnnnnnnnnbbbbbb {inputs}") |
|
res = self.generate_text(inputs) |
|
return res |
|
|
