|
|
| import torch
|
| from transformers import TextIteratorStreamer
|
| from load_model import get_model
|
| import threading
|
|
|
| def generate_response(
|
| user_prompt: str,
|
| system_prompt: str = "You are a helpful AI assistant.",
|
| max_tokens: int = 32768,
|
| stream: bool = False,
|
| ) -> 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()
|
|
|
| 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=True,
|
| temperature=0.8,
|
| top_k=50,
|
| top_p=0.95,
|
| )
|
|
|
| if stream:
|
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
|
| gen_kwargs["streamer"] = streamer
|
| thread = threading.Thread(target=model.generate, kwargs=gen_kwargs)
|
| thread.start()
|
| for text in streamer:
|
| yield text
|
| else:
|
| with torch.no_grad():
|
| outputs = model.generate(**gen_kwargs)
|
|
|
| return tokenizer.decode(
|
| outputs[0][input_ids.shape[1]:],
|
| skip_special_tokens=True,
|
| ).strip()
|
|
|
