wehub-resource-sync
15 lines
748 B
Python
15 lines
748 B
Python
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
from transformers.generation import GenerationConfig
|
|
|
|
|
|
class HFModel(object):
|
|
|
|
def __init__(self, model_path):
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_path,
|
|
trust_remote_code=True,
|
|
device_map='auto',
|
|
low_cpu_mem_usage=True).eval()
|
|
self.model.generation_config = GenerationConfig.from_pretrained(model_path, trust_remote_code=True)
|
|
self.model.generation_config.do_sample = False
|