File size: 2,131 Bytes
7e8f83f | 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 | """Local-HF chat adapter matching the original host defaults."""
from contextlib import redirect_stdout
import json,sys,time
def chat(data):
import torch
from transformers import AutoModelForCausalLM,AutoTokenizer,StoppingCriteria,StoppingCriteriaList
tokenizer=AutoTokenizer.from_pretrained(data['model_path'],local_files_only=True,trust_remote_code=False)
model=AutoModelForCausalLM.from_pretrained(data['model_path'],torch_dtype='auto',device_map={'':'cuda:0'},low_cpu_mem_usage=True,local_files_only=True,trust_remote_code=False)
model.eval()
if data.get('seed') is not None:torch.manual_seed(data['seed'])
prompt=tokenizer.apply_chat_template(data['messages'],tokenize=False,add_generation_prompt=True)
encoded=tokenizer(prompt,return_tensors='pt',truncation=True,max_length=4096).to(model.device)
count=encoded['input_ids'].shape[-1]
deadline=time.monotonic()+max(1,data.get('generation_seconds',250))
class StopAtDeadline(StoppingCriteria):
def __call__(self,input_ids,scores,**kwargs):return time.monotonic()>=deadline
kwargs={'max_new_tokens':data['max_tokens'],'do_sample':data['temperature']>0,'repetition_penalty':data['repetition_penalty'],
'stopping_criteria':StoppingCriteriaList([StopAtDeadline()]),'pad_token_id':tokenizer.pad_token_id or tokenizer.eos_token_id}
if kwargs['do_sample']:kwargs.update(temperature=data['temperature'],top_p=data['top_p'])
with torch.inference_mode():tokens=model.generate(**encoded,**kwargs)
generated=tokens[0,count:]
return {'text':tokenizer.decode(generated,skip_special_tokens=True),'usage':{'prompt_tokens':count,'completion_tokens':len(generated)},
'summary':{'backend':'transformers','finish_reason':'time_limit' if time.monotonic()>=deadline else 'completed'},'files':[]}
if __name__=='__main__':
try:
with redirect_stdout(sys.stderr):result=chat(json.load(sys.stdin))
json.dump(result,sys.stdout,ensure_ascii=False)
except Exception as error:
json.dump({'error':str(error)[:1600]},sys.stdout,ensure_ascii=False)
raise SystemExit(1)
|