48 lines
1.1 KiB
Python
48 lines
1.1 KiB
Python
import os
|
|
|
|
|
|
def make_mock_model_and_tokenizer():
|
|
"""Returns a tuple of HF AutoModelForCausalLM and AutoTokenizer."""
|
|
import torch
|
|
|
|
torch.set_num_threads(1)
|
|
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
# Can regenerate the sample pipe with this:
|
|
# pipeline(
|
|
# "text-generation",
|
|
# "hf-internal-testing/tiny-random-gpt2",
|
|
# ).save_pretrained("...")
|
|
|
|
savedir = os.path.join(
|
|
os.path.abspath(os.path.normpath(os.path.dirname(__file__))), "tiny-random-gpt2"
|
|
)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
savedir,
|
|
local_files_only=True,
|
|
)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
savedir,
|
|
local_files_only=True,
|
|
)
|
|
|
|
return model, tokenizer
|
|
|
|
|
|
def make_mock_pipeline():
|
|
from transformers import pipeline
|
|
|
|
model, tokenizer = make_mock_model_and_tokenizer()
|
|
|
|
pipe = pipeline(
|
|
task="text-generation",
|
|
model=model,
|
|
tokenizer=tokenizer,
|
|
trust_remote_code=False,
|
|
device_map="cpu", # Force CPU to avoid multithreaded fighting.
|
|
)
|
|
return pipe
|