File size: 2,090 Bytes
0a25951 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import argparse
from pathlib import Path
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description="Run the local phi2-merged model.")
parser.add_argument(
"prompt",
nargs="?",
default="Kısa bir selam ver:",
help="Prompt to send to the model.",
)
parser.add_argument(
"--max-new-tokens",
type=int,
default=120,
help="Maximum number of new tokens to generate.",
)
parser.add_argument(
"--temperature",
type=float,
default=0.0,
help="Sampling temperature. Default 0 uses deterministic greedy decoding.",
)
parser.add_argument(
"--top-p",
type=float,
default=0.9,
help="Nucleus sampling threshold.",
)
return parser
def main() -> None:
args = build_parser().parse_args()
model_path = Path(__file__).resolve().parent
tokenizer = AutoTokenizer.from_pretrained(model_path)
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
model = AutoModelForCausalLM.from_pretrained(
model_path,
dtype=dtype,
device_map="auto" if torch.cuda.is_available() else None,
low_cpu_mem_usage=True,
)
inputs = tokenizer(args.prompt, return_tensors="pt")
inputs = {name: tensor.to(model.device) for name, tensor in inputs.items()}
generate_kwargs = {
"max_new_tokens": args.max_new_tokens,
"pad_token_id": tokenizer.eos_token_id,
}
if args.temperature > 0:
generate_kwargs["do_sample"] = True
generate_kwargs["temperature"] = args.temperature
generate_kwargs["top_p"] = args.top_p
else:
generate_kwargs["do_sample"] = False
with torch.no_grad():
output = model.generate(**inputs, **generate_kwargs)
print(tokenizer.decode(output[0], skip_special_tokens=True))
if __name__ == "__main__":
main() |