|
|
|
|
|
import torch
|
|
|
from load_model import get_model
|
|
|
|
|
|
def generate_response(
|
|
|
user_prompt: str,
|
|
|
system_prompt: str = "You are a helpful AI assistant.",
|
|
|
max_tokens: int = 256,
|
|
|
temperature: float = 0.2,
|
|
|
) -> str:
|
|
|
"""Generate response using ALREADY LOADED model"""
|
|
|
model, tokenizer = get_model()
|
|
|
|
|
|
messages = [
|
|
|
{"role": "system", "content": system_prompt},
|
|
|
{"role": "user", "content": user_prompt},
|
|
|
]
|
|
|
|
|
|
input_ids = tokenizer.apply_chat_template(
|
|
|
messages,
|
|
|
add_generation_prompt=True,
|
|
|
return_tensors="pt",
|
|
|
).to(model.device)
|
|
|
|
|
|
attention_mask = (input_ids != tokenizer.pad_token_id).long()
|
|
|
|
|
|
do_sample = temperature > 0.0
|
|
|
gen_kwargs = dict(
|
|
|
input_ids=input_ids,
|
|
|
attention_mask=attention_mask,
|
|
|
max_new_tokens=max_tokens,
|
|
|
pad_token_id=tokenizer.eos_token_id,
|
|
|
use_cache=False,
|
|
|
do_sample=do_sample,
|
|
|
)
|
|
|
if do_sample:
|
|
|
gen_kwargs["temperature"] = float(temperature)
|
|
|
gen_kwargs["top_p"] = 0.95
|
|
|
|
|
|
with torch.no_grad():
|
|
|
outputs = model.generate(**gen_kwargs)
|
|
|
|
|
|
return tokenizer.decode(
|
|
|
outputs[0][input_ids.shape[1]:],
|
|
|
skip_special_tokens=True,
|
|
|
).strip()
|
|
|
|
