File size: 1,387 Bytes
284cb2b |
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 |
import os
import torch
from vllm import LLM, SamplingParams
import logging
# Configure logging
logging.basicConfig(
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)
class ChallengePromptGenerator:
def __init__(
self,
model_local_dir="checkpoint-15000",
):
self.generator = LLM(
model_local_dir,
dtype="bfloat16",
)
def infer_prompt(
self,
prompts,
max_generation_length=77,
beam_size=1,
sampling_temperature=0.9,
sampling_topk=1,
sampling_topp=1,
):
added_prompts = [f"{self.generator.get_tokenizer().bos_token} {prompt}" for prompt in prompts]
sampling_params = SamplingParams(
max_tokens=max_generation_length,
temperature=sampling_temperature,
top_k=sampling_topk,
top_p=sampling_topp,
use_beam_search=(beam_size > 1),
)
outputs = self.generator.generate(added_prompts, sampling_params)
out = []
for i in range(len(outputs)):
tmp_out = prompts[i] + outputs[i].outputs[0].text
if tmp_out[-1] != ".":
tmp_out = ".".join(tmp_out.split(".")[:-1]) + "."
out.append(tmp_out)
return out
|