Spaces:
Sleeping
Sleeping
File size: 1,996 Bytes
eca6215 |
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 |
from transformers import TextStreamer
def prepare_inference_inputs(tokenizer, template, instruction, input_text, eos_token, device="cuda"):
"""
Prepares the inputs for inference by formatting the prompt and tokenizing it.
Args:
- tokenizer: The tokenizer used for tokenization.
- template: The template string for the prompt format.
- instruction: The instruction to be included in the prompt.
- input_text: The input to be included in the prompt.
- eos_token: The end of sequence token.
- device: The device for the model ('cuda' or 'cpu').
Returns:
- Tokenized inputs ready for inference.
"""
prompt = template.format(instruction, input_text, "") + eos_token
return tokenizer([prompt], return_tensors="pt").to(device)
def generate_responses(model, inputs, tokenizer, max_new_tokens=64):
"""
Generates responses from the model based on the provided inputs.
Args:
- model: The pre-trained model for generation.
- inputs: The tokenized inputs to generate responses.
- tokenizer: The tokenizer used to decode the output.
- max_new_tokens: The maximum number of tokens to generate.
Returns:
- Decoded responses from the model.
"""
outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, use_cache=True)
return tokenizer.batch_decode(outputs, skip_special_tokens=True)
def stream_responses(model, inputs, tokenizer, max_new_tokens=128):
"""
Streams the model's response using a text streamer.
Args:
- model: The pre-trained model for generation.
- inputs: The tokenized inputs to generate responses.
- tokenizer: The tokenizer used to decode the output.
- max_new_tokens: The maximum number of tokens to generate.
Returns:
- Streams the output directly.
"""
text_streamer = TextStreamer(tokenizer)
model.generate(**inputs, streamer=text_streamer, max_new_tokens=max_new_tokens)
|