File size: 4,155 Bytes
0b7ba67 8483978 0b7ba67 191c0de 8483978 191c0de 8483978 191c0de 0b7ba67 191c0de 0b7ba67 191c0de 0b7ba67 191c0de 0b7ba67 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
"""
Utility functions for pre-processing and post-processing in the chat application.
"""
def preprocess_chat_input(message, history, system_prompt=""):
"""
Pre-process chat input to prepare it for the model.
Args:
message (str): The current user message
history (list): List of tuples containing (user_message, assistant_message) pairs
system_prompt (str): Optional system prompt to guide the model's behavior
Returns:
dict: Formatted messages in the format expected by the tokenizer
"""
# Convert history from tuples to dict format expected by apply_chat_template
formatted_history = []
for user_msg, assistant_msg in history:
formatted_history.append({"role": "user", "content": user_msg})
formatted_history.append({"role": "assistant", "content": assistant_msg})
# Add current message
formatted_history.append({"role": "user", "content": message})
# Add system message if provided
if system_prompt.strip():
messages = [{"role": "system", "content": system_prompt.strip()}] + formatted_history
else:
messages = formatted_history
return messages
def format_prompt(message, history, tokenizer, system_prompt=""):
"""
Format message and history into a prompt for Qwen models.
Uses tokenizer.apply_chat_template if available, otherwise falls back to manual formatting.
Args:
message (str): The current user message
history (list): List of tuples containing (user_message, assistant_message) pairs
tokenizer: The model tokenizer
system_prompt (str): Optional system prompt to guide the model's behavior
Returns:
str: Formatted prompt ready for the model
"""
# Get pre-processed messages
messages = preprocess_chat_input(message, history, system_prompt)
# Apply chat template if available
if hasattr(tokenizer, "chat_template") and tokenizer.chat_template:
return tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
enable_thinking=True
)
else:
# Fallback for base LMs without chat template
return format_prompt_fallback(messages)
def format_prompt_fallback(messages):
"""
Fallback prompt formatting for models without chat templates.
Args:
messages (list): List of message dictionaries with role and content
Returns:
str: Formatted prompt string
"""
prompt = ""
# Add system message if present
if messages and messages[0]['role'] == 'system':
prompt = messages[0]['content'].strip() + "\n"
messages = messages[1:]
# Add conversation history
for msg in messages:
if msg['role'] == 'user':
prompt += f"<|User|>: {msg['content'].strip()}\n"
elif msg['role'] == 'assistant':
prompt += f"<|Assistant|>: {msg['content'].strip()}\n"
# Add final assistant prompt if needed
if not prompt.strip().endswith("<|Assistant|>:"):
prompt += "<|Assistant|>:"
return prompt
def prepare_generation_inputs(prompt, tokenizer, device):
"""
Prepare tokenized inputs for model generation.
Args:
prompt (str): The formatted prompt
tokenizer: The model tokenizer
device: The device to place tensors on
Returns:
dict: Tokenized inputs ready for model generation
"""
# Tokenize the prompt
tokenized_inputs = tokenizer(prompt, return_tensors="pt")
# Move tensors to the correct device
inputs_on_device = {k: v.to(device) for k, v in tokenized_inputs.items()}
return inputs_on_device
def postprocess_response(response):
"""
Post-process the model's response.
Args:
response (str): The raw model response
Returns:
str: The processed response
"""
# Currently just returns the response as-is
# This function can be expanded for additional post-processing steps
return response |