+ # blocks and any text between them (thinking blocks, intermediate text).
+ if tool_func_names and not visible_prefix and _model_visible:
+ tc_matches = list(re.finditer(r'.*?', _model_visible, re.DOTALL))
+ if tc_matches:
+ prefix_end = tc_matches[-1].end()
+ prefix = _model_visible[:prefix_end].strip()
+ if prefix:
+ visible_prefix = [prefix]
+ _model_visible = _model_visible[prefix_end:].strip()
+
+ # Re-apply visible prefix to the final state after streaming completes.
+ # This is safe because we're no longer sharing the object with chatbot_wrapper.
+ if visible_prefix:
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix + [_model_visible])
+
+ if tool_func_names:
save_history(history, state['unique_id'], state['character_menu'], state['mode'])
- last_save_time = current_time
+
+ # Check for tool calls
+ if not tool_func_names or shared.stop_everything:
+ break
+
+ answer = history['internal'][-1][1]
+ parsed_calls, content_prefix = parse_tool_call(answer, tool_func_names, return_prefix=True, parsers=_tool_parsers) if answer else (None, '')
+
+ if not parsed_calls:
+ break # No tool calls — done
+
+ # --- Process tool calls ---
+ row_idx = len(history['internal']) - 1
+ meta = history.get('metadata', {})
+ seq = meta.setdefault(f'assistant_{row_idx}', {}).setdefault('tool_sequence', [])
+
+ def _render():
+ return chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+
+ # Serialize tool calls and build display headers in one pass
+ serialized = []
+ tc_headers = []
+ for tc in parsed_calls:
+ tc['id'] = get_tool_call_id()
+ fn_name = tc['function']['name']
+ fn_args = tc['function'].get('arguments', {})
+
+ serialized.append({
+ 'id': tc['id'],
+ 'type': 'function',
+ 'function': {
+ 'name': fn_name,
+ 'arguments': json.dumps(fn_args) if isinstance(fn_args, dict) else fn_args
+ }
+ })
+
+ if isinstance(fn_args, dict) and fn_args:
+ args_summary = ', '.join(f'{k}={json.dumps(v, ensure_ascii=False)}' for k, v in fn_args.items())
+ elif isinstance(fn_args, dict):
+ args_summary = ''
+ else:
+ args_summary = str(fn_args)
+
+ tc_headers.append(f'{fn_name}({args_summary})')
+
+ seq_entry = {'tool_calls': serialized}
+ if content_prefix.strip():
+ # Strip GPT-OSS channel tokens so they don't get double-wrapped
+ # by the template (which adds its own channel markup).
+ clean = content_prefix.strip()
+ if '<|channel|>' in clean and '<|message|>' in clean:
+ inner = clean.split('<|message|>', 1)[1]
+ if '<|end|>' in inner:
+ inner = inner.split('<|end|>', 1)[0]
+ clean = inner.strip()
+ if clean:
+ seq_entry['content'] = clean
+ seq.append(seq_entry)
+
+ # Clear internal (raw tool markup)
+ history['internal'][-1][1] = ''
+
+ # Preserve thinking block and intermediate text from this turn.
+ # content_prefix is the raw text before tool call syntax (returned
+ # by parse_tool_call); HTML-escape it and extract thinking to get
+ # the content the user should see.
+ content_text = html.escape(content_prefix)
+ thinking_content, intermediate = extract_thinking_block(content_text)
+ if thinking_content:
+ visible_prefix.append(f'<think>\n{thinking_content}\n</think>')
+ if intermediate and intermediate.strip():
+ visible_prefix.append(intermediate.strip())
+
+ # Show placeholder accordions with "..." before execution starts
+ # (tool calls may be slow, e.g. web search).
+ pending_placeholders = [f'{h}\n...\n' for h in tc_headers]
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix + pending_placeholders)
+ yield _render(), history
+
+ # Execute tools, store results, and replace placeholders with real results
+ for i, tc in enumerate(parsed_calls):
+ # Check for stop request before each tool execution
+ if shared.stop_everything:
+ for j in range(i, len(parsed_calls)):
+ seq.append({'role': 'tool', 'content': 'Tool execution was cancelled by the user.', 'tool_call_id': parsed_calls[j]['id']})
+ pending_placeholders[j] = f'{tc_headers[j]}\nCancelled\n'
+
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix + pending_placeholders)
+ yield _render(), history
+ break
+
+ fn_name = tc['function']['name']
+ fn_args = tc['function'].get('arguments', {})
+ result = execute_tool(fn_name, fn_args, tool_executors)
+
+ seq.append({'role': 'tool', 'content': result, 'tool_call_id': tc['id']})
+ try:
+ pretty_result = json.dumps(json.loads(result), indent=2, ensure_ascii=False)
+ except (json.JSONDecodeError, TypeError):
+ pretty_result = result
+
+ # Replace the placeholder with the real result
+ pending_placeholders[i] = f'{tc_headers[i]}\n{pretty_result}\n'
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix + pending_placeholders)
+ yield _render(), history
+
+ # Move completed tool calls into visible_prefix for next turns
+ visible_prefix.extend(pending_placeholders)
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix)
+ save_history(history, state['unique_id'], state['character_menu'], state['mode'])
+
+ state['history'] = history
+ _tool_turn += 1
+
+ state.pop('_tool_turn', None)
+
+ # If output extensions were deferred during tool turns, apply them now
+ # to the final model response only (not to tool call markers).
+ if state.pop('_skip_output_extensions', None):
+ _model_visible = apply_extensions('output', _model_visible, state, is_chat=True)
+ if visible_prefix:
+ history['visible'][-1][1] = '\n\n'.join(visible_prefix + [_model_visible])
+ else:
+ history['visible'][-1][1] = _model_visible
+
+ yield chat_html_wrapper(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu']), history
+
+ state['history'] = history
+
+ # Sync version metadata so swipes show the full visible (with tool prefix)
+ if visible_prefix and history.get('metadata'):
+ row_idx = len(history['internal']) - 1
+ key = f"assistant_{row_idx}"
+ meta_entry = history['metadata'].get(key, {})
+ if 'versions' in meta_entry and 'current_version_index' in meta_entry:
+ current_idx = meta_entry['current_version_index']
+ if current_idx < len(meta_entry['versions']):
+ version_update = {
+ 'content': history['internal'][row_idx][1],
+ 'visible_content': history['visible'][row_idx][1]
+ }
+ ts = meta_entry.get('tool_sequence')
+ if ts is not None:
+ version_update['tool_sequence'] = ts
+ meta_entry['versions'][current_idx].update(version_update)
save_history(history, state['unique_id'], state['character_menu'], state['mode'])
@@ -1164,7 +1573,7 @@ def redraw_html(history, name1, name2, mode, style, character, reset_cache=False
return chat_html_wrapper(history, name1, name2, mode, style, character, reset_cache=reset_cache)
-def start_new_chat(state):
+def start_new_chat(state, unique_id=None):
mode = state['mode']
# Initialize with empty metadata dictionary
history = {'internal': [], 'visible': [], 'metadata': {}}
@@ -1178,7 +1587,9 @@ def start_new_chat(state):
# Add timestamp for assistant's greeting
update_message_metadata(history['metadata'], "assistant", 0, timestamp=get_current_timestamp())
- unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')
+ if unique_id is None:
+ unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')
+
save_history(history, unique_id, state['character_menu'], state['mode'])
return history
@@ -1197,12 +1608,16 @@ def save_history(history, unique_id, character, mode):
if shared.args.multi_user:
return
+ if unique_id and unique_id.startswith('incognito-'):
+ return
+
p = get_history_file_path(unique_id, character, mode)
if not p.parent.is_dir():
p.parent.mkdir(parents=True)
- with open(p, 'w', encoding='utf-8') as f:
- f.write(json.dumps(history, indent=4, ensure_ascii=False))
+ with _history_file_lock:
+ with open(p, 'w', encoding='utf-8') as f:
+ f.write(json.dumps(history, indent=4, ensure_ascii=False))
def rename_history(old_id, new_id, character, mode):
@@ -1333,6 +1748,7 @@ def load_history_after_deletion(state, idx):
Loads the latest history for the given character in chat or chat-instruct
mode, or the latest instruct history for instruct mode.
'''
+ import gradio as gr
if shared.args.multi_user:
return start_new_chat(state)
@@ -1351,6 +1767,7 @@ def load_history_after_deletion(state, idx):
def update_character_menu_after_deletion(idx):
+ import gradio as gr
characters = utils.get_available_characters()
idx = min(int(idx), len(characters) - 1)
idx = max(0, idx)
@@ -1383,6 +1800,9 @@ def save_last_chat_state(character, mode, unique_id):
if shared.args.multi_user:
return
+ if unique_id and unique_id.startswith('incognito-'):
+ return
+
state = load_last_chat_state()
key = get_chat_state_key(character, mode)
state["last_chats"][key] = unique_id
@@ -1565,24 +1985,6 @@ def clear_character_for_ui(state):
return state, state['name2'], state['context'], state['greeting'], None
-def load_instruction_template(template):
- if template == 'None':
- return ''
-
- for filepath in [shared.user_data_dir / 'instruction-templates' / f'{template}.yaml', shared.user_data_dir / 'instruction-templates' / 'Alpaca.yaml']:
- if filepath.exists():
- break
- else:
- return ''
-
- file_contents = open(filepath, 'r', encoding='utf-8').read()
- data = yaml.safe_load(file_contents)
- if 'instruction_template' in data:
- return data['instruction_template']
- else:
- return jinja_template_from_old_format(data)
-
-
@functools.cache
def load_character_memoized(character, name1, name2):
return load_character(character, name1, name2)
@@ -1590,10 +1992,12 @@ def load_character_memoized(character, name1, name2):
@functools.cache
def load_instruction_template_memoized(template):
+ from modules.models_settings import load_instruction_template
return load_instruction_template(template)
def upload_character(file, img_path, tavern=False):
+ import gradio as gr
img = open_image_safely(img_path)
decoded_file = file if isinstance(file, str) else file.decode('utf-8')
try:
@@ -1647,6 +2051,7 @@ def upload_tavern_character(img_path, _json):
def check_tavern_character(img_path):
+ import gradio as gr
img = open_image_safely(img_path)
if img is None:
@@ -1832,6 +2237,7 @@ def delete_user(name):
def update_user_menu_after_deletion(idx):
"""Update user menu after a user is deleted"""
+ import gradio as gr
users = get_available_users()
if len(users) == 0:
# Create a default user if none exist
@@ -1864,93 +2270,13 @@ def handle_user_menu_change(state):
def handle_save_user_click(name1):
"""Handle save user button click"""
+ import gradio as gr
return [
name1,
gr.update(visible=True)
]
-def jinja_template_from_old_format(params, verbose=False):
- MASTER_TEMPLATE = """
-{%- set ns = namespace(found=false) -%}
-{%- for message in messages -%}
- {%- if message['role'] == 'system' -%}
- {%- set ns.found = true -%}
- {%- endif -%}
-{%- endfor -%}
-{%- if not ns.found -%}
- {{- '<|PRE-SYSTEM|>' + '<|SYSTEM-MESSAGE|>' + '<|POST-SYSTEM|>' -}}
-{%- endif %}
-{%- for message in messages %}
- {%- if message['role'] == 'system' -%}
- {{- '<|PRE-SYSTEM|>' + message['content'] + '<|POST-SYSTEM|>' -}}
- {%- else -%}
- {%- if message['role'] == 'user' -%}
- {{-'<|PRE-USER|>' + message['content'] + '<|POST-USER|>'-}}
- {%- else -%}
- {{-'<|PRE-ASSISTANT|>' + message['content'] + '<|POST-ASSISTANT|>' -}}
- {%- endif -%}
- {%- endif -%}
-{%- endfor -%}
-{%- if add_generation_prompt -%}
- {{-'<|PRE-ASSISTANT-GENERATE|>'-}}
-{%- endif -%}
-"""
-
- if 'context' in params and '<|system-message|>' in params['context']:
- pre_system = params['context'].split('<|system-message|>')[0]
- post_system = params['context'].split('<|system-message|>')[1]
- else:
- pre_system = ''
- post_system = ''
-
- pre_user = params['turn_template'].split('<|user-message|>')[0].replace('<|user|>', params['user'])
- post_user = params['turn_template'].split('<|user-message|>')[1].split('<|bot|>')[0]
-
- pre_assistant = '<|bot|>' + params['turn_template'].split('<|bot-message|>')[0].split('<|bot|>')[1]
- pre_assistant = pre_assistant.replace('<|bot|>', params['bot'])
- post_assistant = params['turn_template'].split('<|bot-message|>')[1]
-
- def preprocess(string):
- return string.replace('\n', '\\n').replace('\'', '\\\'')
-
- pre_system = preprocess(pre_system)
- post_system = preprocess(post_system)
- pre_user = preprocess(pre_user)
- post_user = preprocess(post_user)
- pre_assistant = preprocess(pre_assistant)
- post_assistant = preprocess(post_assistant)
-
- if verbose:
- print(
- '\n',
- repr(pre_system) + '\n',
- repr(post_system) + '\n',
- repr(pre_user) + '\n',
- repr(post_user) + '\n',
- repr(pre_assistant) + '\n',
- repr(post_assistant) + '\n',
- )
-
- result = MASTER_TEMPLATE
- if 'system_message' in params:
- result = result.replace('<|SYSTEM-MESSAGE|>', preprocess(params['system_message']))
- else:
- result = result.replace('<|SYSTEM-MESSAGE|>', '')
-
- result = result.replace('<|PRE-SYSTEM|>', pre_system)
- result = result.replace('<|POST-SYSTEM|>', post_system)
- result = result.replace('<|PRE-USER|>', pre_user)
- result = result.replace('<|POST-USER|>', post_user)
- result = result.replace('<|PRE-ASSISTANT|>', pre_assistant)
- result = result.replace('<|PRE-ASSISTANT-GENERATE|>', pre_assistant.rstrip(' '))
- result = result.replace('<|POST-ASSISTANT|>', post_assistant)
-
- result = result.strip()
-
- return result
-
-
def my_yaml_output(data):
'''
pyyaml is very inconsistent with multiline strings.
@@ -2002,6 +2328,7 @@ def handle_unique_id_select(state):
def handle_start_new_chat_click(state):
+ import gradio as gr
history = start_new_chat(state)
histories = find_all_histories_with_first_prompts(state)
html = redraw_html(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
@@ -2016,10 +2343,29 @@ def handle_start_new_chat_click(state):
return [history, html, past_chats_update]
+def handle_start_incognito_chat_click(state):
+ import gradio as gr
+ unique_id = 'incognito-' + datetime.now().strftime('%Y%m%d-%H-%M-%S')
+ history = start_new_chat(state, unique_id=unique_id)
+ html = redraw_html(history, state['name1'], state['name2'], state['mode'], state['chat_style'], state['character_menu'])
+
+ convert_to_markdown.cache_clear()
+
+ histories = find_all_histories_with_first_prompts(state)
+ past_chats_update = gr.update(choices=histories, value=unique_id)
+
+ return [history, html, past_chats_update]
+
+
def handle_delete_chat_confirm_click(state):
filtered_histories = find_all_histories_with_first_prompts(state)
filtered_ids = [h[1] for h in filtered_histories]
- index = str(filtered_ids.index(state['unique_id']))
+
+ if state['unique_id'] not in filtered_ids:
+ # Incognito or unknown chat — just load the most recent saved chat
+ index = '0'
+ else:
+ index = str(filtered_ids.index(state['unique_id']))
delete_history(state['unique_id'], state['character_menu'], state['mode'])
history, unique_id = load_history_after_deletion(state, index)
@@ -2027,16 +2373,11 @@ def handle_delete_chat_confirm_click(state):
convert_to_markdown.cache_clear()
- return [
- history,
- html,
- unique_id,
- gr.update(visible=False),
- gr.update(visible=True),
- ]
+ return [history, html, unique_id]
def handle_branch_chat_click(state):
+ import gradio as gr
branch_from_index = state['branch_index']
if branch_from_index == -1:
history = state['history']
@@ -2048,7 +2389,8 @@ def handle_branch_chat_click(state):
if 'metadata' in history:
history['metadata'] = {k: v for k, v in history['metadata'].items() if int(k.split('_')[-1]) <= branch_from_index}
- new_unique_id = datetime.now().strftime('%Y%m%d-%H-%M-%S')
+ prefix = 'incognito-' if state['unique_id'] and state['unique_id'].startswith('incognito-') else ''
+ new_unique_id = prefix + datetime.now().strftime('%Y%m%d-%H-%M-%S')
save_history(history, new_unique_id, state['character_menu'], state['mode'])
histories = find_all_histories_with_first_prompts(state)
@@ -2086,14 +2428,19 @@ def handle_edit_message_click(state):
original_visible = history['visible'][message_index][role_idx]
original_timestamp = history['metadata'][key].get('timestamp', get_current_timestamp())
- history['metadata'][key]["versions"] = [{
+ version_entry = {
"content": original_content,
"visible_content": original_visible,
"timestamp": original_timestamp
- }]
+ }
+ ts = history['metadata'][key].get('tool_sequence')
+ if ts is not None:
+ version_entry['tool_sequence'] = ts
+ history['metadata'][key]["versions"] = [version_entry]
history['internal'][message_index][role_idx] = apply_extensions('input', new_text, state, is_chat=True)
history['visible'][message_index][role_idx] = html.escape(new_text)
+ history['metadata'][key].pop('tool_sequence', None)
add_message_version(history, role, message_index, is_current=True)
@@ -2138,6 +2485,14 @@ def handle_navigate_version_click(state):
history['internal'][message_index][msg_content_idx] = version_to_load['content']
history['visible'][message_index][msg_content_idx] = version_to_load['visible_content']
metadata['current_version_index'] = new_idx
+
+ # Restore per-version tool_sequence so follow-up prompts see consistent context
+ version_ts = version_to_load.get('tool_sequence')
+ if version_ts is not None:
+ metadata['tool_sequence'] = version_ts
+ else:
+ metadata.pop('tool_sequence', None)
+
update_message_metadata(history['metadata'], role, message_index, timestamp=version_to_load['timestamp'])
# Redraw and save
@@ -2148,6 +2503,7 @@ def handle_navigate_version_click(state):
def handle_rename_chat_click():
+ import gradio as gr
return [
gr.update(value="My New Chat"),
gr.update(visible=True),
@@ -2155,6 +2511,14 @@ def handle_rename_chat_click():
def handle_rename_chat_confirm(rename_to, state):
+ import gradio as gr
+
+ if state['unique_id'] and state['unique_id'].startswith('incognito-'):
+ return [
+ gr.update(),
+ gr.update(visible=False),
+ ]
+
rename_history(state['unique_id'], rename_to, state['character_menu'], state['mode'])
histories = find_all_histories_with_first_prompts(state)
@@ -2165,11 +2529,13 @@ def handle_rename_chat_confirm(rename_to, state):
def handle_search_chat_change(state):
+ import gradio as gr
histories = find_all_histories_with_first_prompts(state)
return gr.update(choices=histories)
def handle_upload_chat_history(load_chat_history, state):
+ import gradio as gr
history = start_new_chat(state)
history = load_history_json(load_chat_history, history)
save_history(history, state['unique_id'], state['character_menu'], state['mode'])
@@ -2192,6 +2558,7 @@ def handle_upload_chat_history(load_chat_history, state):
def handle_character_menu_change(state):
+ import gradio as gr
name1, name2, picture, greeting, context = load_character(state['character_menu'], state['name1'], state['name2'])
state['name1'] = name1
@@ -2244,6 +2611,7 @@ def handle_character_picture_change(picture_path):
def handle_mode_change(state):
+ import gradio as gr
history, loaded_unique_id = load_latest_history(state)
histories = find_all_histories_with_first_prompts(state)
@@ -2270,6 +2638,7 @@ def handle_mode_change(state):
def handle_save_character_click(name2):
+ import gradio as gr
return [
name2,
gr.update(visible=True)
@@ -2277,6 +2646,7 @@ def handle_save_character_click(name2):
def handle_load_template_click(instruction_template):
+ from modules.models_settings import load_instruction_template
output = load_instruction_template(instruction_template)
return [
output,
@@ -2285,20 +2655,26 @@ def handle_load_template_click(instruction_template):
def handle_save_template_click(instruction_template_str):
+ import gradio as gr
contents = generate_instruction_template_yaml(instruction_template_str)
+ root = str(shared.user_data_dir / 'instruction-templates') + '/'
return [
"My Template.yaml",
- str(shared.user_data_dir / 'instruction-templates') + '/',
+ root,
contents,
+ root,
gr.update(visible=True)
]
def handle_delete_template_click(template):
+ import gradio as gr
+ root = str(shared.user_data_dir / 'instruction-templates') + '/'
return [
f"{template}.yaml",
- str(shared.user_data_dir / 'instruction-templates') + '/',
- gr.update(visible=False)
+ root,
+ root,
+ gr.update(visible=True)
]
@@ -2310,6 +2686,7 @@ def handle_your_picture_change(picture, state):
def handle_send_instruction_click(state):
+ import gradio as gr
state['mode'] = 'instruct'
state['history'] = {'internal': [], 'visible': [], 'metadata': {}}
@@ -2322,6 +2699,7 @@ def handle_send_instruction_click(state):
def handle_send_chat_click(state):
+ import gradio as gr
output = generate_chat_prompt("", state, _continue=True)
if state["show_two_notebook_columns"]:
diff --git a/modules/exllamav3.py b/modules/exllamav3.py
index b4b76e21..e1efbfeb 100644
--- a/modules/exllamav3.py
+++ b/modules/exllamav3.py
@@ -1,6 +1,6 @@
+import math
import queue
import threading
-import traceback
from pathlib import Path
from typing import Any, List, Tuple
@@ -9,6 +9,7 @@ import torch
from exllamav3 import Cache, Config, Generator, Model, Tokenizer
from exllamav3.cache import CacheLayer_fp16, CacheLayer_quant
from exllamav3.generator import Job
+from exllamav3.generator.filter import Filter
from exllamav3.generator.sampler import (
CustomSampler,
SS_AdaptiveP,
@@ -32,8 +33,30 @@ from modules.text_generation import get_max_prompt_length
try:
import flash_attn
except Exception:
- logger.warning('Failed to load flash-attention due to the following error:\n')
- traceback.print_exc()
+ logger.warning('Failed to load flash-attention due to the following error:', exc_info=True)
+
+
+class LogitBiasFilter(Filter):
+ """Filter subclass that applies a static additive logit bias mask."""
+
+ def __init__(self, tokenizer, logit_bias_dict):
+ super().__init__(tokenizer=tokenizer, trigger_token=None, prefix_str=None, eos_after_completed=False)
+ self.logit_bias_dict = logit_bias_dict
+ self._mask = None
+
+ def reset(self): pass
+ def accept_token(self, token): pass
+ def is_completed(self): return False
+ def use_background_worker(self): return False
+
+ def get_next_logit_mask(self):
+ if self._mask is None:
+ self._mask = torch.zeros((1, self.vocab_size), dtype=self.logits_dtype)
+ for token_id_str, bias in self.logit_bias_dict.items():
+ token_id = int(token_id_str)
+ if 0 <= token_id < self.vocab_size:
+ self._mask[0, token_id] = bias
+ return self._mask
class ConcurrentGenerator:
@@ -53,7 +76,16 @@ class ConcurrentGenerator:
if not self.job_queues:
self.has_jobs.clear()
continue
- results = self.generator.iterate()
+ try:
+ results = self.generator.iterate()
+ except Exception:
+ logger.exception("Exception in ConcurrentGenerator iterate loop")
+ for q in self.job_queues.values():
+ q.put(None)
+ self.job_queues.clear()
+ self.generator.clear_queue()
+ self.has_jobs.clear()
+ continue
for result in results:
job = result["job"]
q = self.job_queues.get(job)
@@ -89,6 +121,10 @@ class Exllamav3Model:
def __init__(self):
pass
+ @property
+ def device(self) -> torch.device:
+ return torch.device(0)
+
@classmethod
def from_pretrained(cls, path_to_model):
path_to_model = Path(f'{shared.args.model_dir}') / Path(path_to_model)
@@ -149,8 +185,21 @@ class Exllamav3Model:
load_params['tensor_p'] = True
load_params['tp_backend'] = shared.args.tp_backend
- model.load(**load_params)
- tokenizer = Tokenizer.from_config(config)
+ # Load vision and draft before the main model so autosplit
+ # accounts for their VRAM usage.
+
+ # Load vision model component (ExLlamaV3 native)
+ vision_model = None
+ if "vision_config" in config.config_dict:
+ logger.info("Vision component detected in model config. Attempting to load...")
+ try:
+ vision_model = Model.from_config(config, component="vision")
+ vision_model.load(progressbar=True)
+ logger.info("Vision model loaded successfully.")
+ except Exception as e:
+ logger.warning(f"Vision model loading failed (multimodal disabled): {e}")
+ else:
+ logger.info("No vision component in model config. Skipping multimodal setup.")
# Initialize draft model for speculative decoding
draft_model = None
@@ -166,23 +215,8 @@ class Exllamav3Model:
logger.warning(f"Draft model not found at {draft_path}, speculative decoding disabled.")
else:
draft_config = Config.from_directory(str(draft_path))
-
- # Set context size for draft model with 256-multiple validation
- if shared.args.ctx_size_draft > 0:
- draft_max_tokens = shared.args.ctx_size_draft
- else:
- draft_max_tokens = shared.args.ctx_size
-
- # Validate draft model context size is a multiple of 256
- if draft_max_tokens % 256 != 0:
- adjusted_draft_tokens = ((draft_max_tokens // 256) + 1) * 256
- logger.warning(f"Draft model max_num_tokens must be a multiple of 256. Adjusting from {draft_max_tokens} to {adjusted_draft_tokens}")
- draft_max_tokens = adjusted_draft_tokens
-
- draft_config.max_seq_len = draft_max_tokens
-
draft_model = Model.from_config(draft_config)
- draft_cache = Cache(draft_model, max_num_tokens=draft_max_tokens, layer_type=layer_type, **cache_kwargs)
+ draft_cache = Cache(draft_model, max_num_tokens=max_tokens, layer_type=layer_type, **cache_kwargs)
draft_load_params = {'progressbar': True}
if split:
@@ -191,18 +225,9 @@ class Exllamav3Model:
draft_model.load(**draft_load_params)
logger.info(f"Draft model loaded successfully. Max speculative tokens: {shared.args.draft_max}")
- # Load vision model component (ExLlamaV3 native)
- vision_model = None
- if "vision_config" in config.config_dict:
- logger.info("Vision component detected in model config. Attempting to load...")
- try:
- vision_model = Model.from_config(config, component="vision")
- vision_model.load(progressbar=True)
- logger.info("Vision model loaded successfully.")
- except Exception as e:
- logger.warning(f"Vision model loading failed (multimodal disabled): {e}")
- else:
- logger.info("No vision component in model config. Skipping multimodal setup.")
+ # Load main model last
+ model.load(**load_params)
+ tokenizer = Tokenizer.from_config(config)
generator = Generator(
model=model,
@@ -385,11 +410,31 @@ class Exllamav3Model:
else:
max_new_tokens = state['max_new_tokens']
- # Get stop conditions
+ # Use full EOS token list from config (may contain multiple IDs)
stop_conditions = []
if not state['ban_eos_token']:
- if hasattr(self.tokenizer, 'eos_token_id') and self.tokenizer.eos_token_id is not None:
- stop_conditions.append(self.tokenizer.eos_token_id)
+ for eos_id in self.config.eos_token_id_list:
+ if eos_id is not None:
+ stop_conditions.append(eos_id)
+
+ # Build filters for logit_bias (OpenAI API)
+ filters = []
+ logit_bias = state.get('logit_bias')
+ if logit_bias:
+ filters.append(LogitBiasFilter(self.tokenizer, logit_bias))
+
+ # Suppress EOS tokens via logit bias so they are never sampled
+ if state['ban_eos_token']:
+ eos_bias = {}
+ for eos_id in self.config.eos_token_id_list:
+ if eos_id is not None:
+ eos_bias[str(eos_id)] = float('-inf')
+ if eos_bias:
+ filters.append(LogitBiasFilter(self.tokenizer, eos_bias))
+
+ # Logprobs support (OpenAI API)
+ logprobs = state.get('logprobs', 0) or 0
+ return_top_tokens = logprobs if logprobs > 0 else 0
seed = state.get('seed', -1)
job = Job(
@@ -400,11 +445,15 @@ class Exllamav3Model:
sampler=sampler,
seed=seed if seed >= 0 else None,
stop_conditions=stop_conditions if stop_conditions else None,
+ filters=filters if filters else None,
+ return_top_tokens=return_top_tokens,
+ return_probs=return_top_tokens > 0,
)
# Stream generation
response_text = ""
stop_event = state.get('stop_event')
+ self.last_completion_probabilities = []
result_queue = self.parallel_generator.submit(job)
try:
@@ -416,14 +465,61 @@ class Exllamav3Model:
except queue.Empty:
continue
if result is None or result.get("eos"):
+ # Capture logprobs from the final eos result too
+ if result is not None and return_top_tokens > 0:
+ self._capture_logprobs(result)
break
chunk = result.get("text", "")
+
+ # Capture logprobs from streaming results
+ if return_top_tokens > 0:
+ self._capture_logprobs(result)
+
if chunk:
response_text += chunk
yield response_text
finally:
self.parallel_generator.cancel(job)
+ def _capture_logprobs(self, result):
+ """Convert ExLlamav3 top-k token data to the shared logprobs format."""
+ top_k_tokens = result.get("top_k_tokens")
+ top_k_probs = result.get("top_k_probs")
+ if top_k_tokens is None or top_k_probs is None:
+ return
+
+ id_to_piece = self.tokenizer.get_id_to_piece_list(True)
+ sampled_ids = result.get("token_ids") # (batch, seq_len) - actually sampled tokens
+ sampled_probs = result.get("token_probs") # (batch, seq_len) - their probabilities
+
+ def _piece(tid):
+ s = id_to_piece[tid] if tid < len(id_to_piece) else f"<{tid}>"
+ return s.replace('\u2581', ' ')
+
+ def _logprob(prob):
+ return math.log(prob) if prob > 0 else float("-inf")
+
+ # top_k_tokens shape: (batch, seq_len, k), top_k_probs same
+ for seq_idx in range(top_k_tokens.shape[1]):
+ entry = {"top_logprobs": []}
+ for k_idx in range(top_k_tokens.shape[2]):
+ token_id = top_k_tokens[0, seq_idx, k_idx].item()
+ prob = top_k_probs[0, seq_idx, k_idx].item()
+ entry["top_logprobs"].append({"token": _piece(token_id), "logprob": _logprob(prob)})
+
+ # Record the actually sampled token at the entry level so
+ # format_completion_logprobs uses it instead of top_logprobs[0]
+ # (they differ with non-greedy sampling).
+ if sampled_ids is not None:
+ sid = sampled_ids[0, seq_idx].item()
+ entry["token"] = _piece(sid)
+ if sampled_probs is not None:
+ entry["logprob"] = _logprob(sampled_probs[0, seq_idx].item())
+ else:
+ entry["logprob"] = None
+
+ self.last_completion_probabilities.append(entry)
+
def generate(self, prompt, state):
output = ""
for chunk in self.generate_with_streaming(prompt, state):
@@ -431,42 +527,31 @@ class Exllamav3Model:
return output
+ def get_prompt_logits(self, input_ids):
+ """Return logits for all positions via a single no-cache forward pass.
+
+ Used by prompt logprobs computation. Returns (1, seq_len, vocab) on CPU in float32.
+ """
+ import torch
+ input_ids_tensor = input_ids if isinstance(input_ids, torch.Tensor) else torch.tensor(input_ids, dtype=torch.long)
+ input_ids_tensor = input_ids_tensor.view(1, -1).cpu()
+ with torch.no_grad():
+ return self.model.forward(
+ input_ids=input_ids_tensor,
+ params={"attn_mode": "flash_attn_nc"}
+ ).cpu().float()
+
def get_logits(self, token_ids, **kwargs):
"""
Process a batch of token_ids and return the logits for the last token.
- This will reset and overwrite the model's cache.
+ Uses flash_attn_nc (no cache) for correct results with recurrent models.
"""
- # Initialize a single params dictionary that will be updated in-place
- params = {
- "cache": self.cache,
- "reconstruct": False,
- "attn_mode": "flash_attn",
- "batch_shape": (1, self.max_tokens),
- "past_len": 0
- }
- params.update(kwargs)
-
- # Process prefix tokens to fill the cache and generate recurrent state
- if token_ids.shape[-1] > 1:
- prefix_ids = token_ids[:, :-1]
-
- # This forward call updates the 'params' dict with the recurrent state
- self.model.forward(
- input_ids=prefix_ids,
- params=params
- )
-
- # Update past_len for the next call
- params["past_len"] = prefix_ids.shape[-1]
-
- # Process the last token, now using the state-filled 'params' dict
- last_token_ids = token_ids[:, -1:]
logits = self.model.forward(
- input_ids=last_token_ids,
- params=params
+ input_ids=token_ids,
+ params={"attn_mode": "flash_attn_nc"}
)
- return logits.float().cpu()
+ return logits[:, -1:, :].float().cpu()
def encode(self, string, **kwargs):
add_bos = kwargs.pop('add_bos', True)
diff --git a/modules/exllamav3_hf.py b/modules/exllamav3_hf.py
index b4b6ad20..5e634e22 100644
--- a/modules/exllamav3_hf.py
+++ b/modules/exllamav3_hf.py
@@ -1,5 +1,4 @@
import os
-import traceback
from pathlib import Path
from typing import Any, Dict, Optional, Union
@@ -21,13 +20,15 @@ from modules.logging_colors import logger
try:
import flash_attn
except Exception:
- logger.warning('Failed to load flash-attention due to the following error:\n')
- traceback.print_exc()
+ logger.warning('Failed to load flash-attention due to the following error:', exc_info=True)
class Exllamav3HF(PreTrainedModel, GenerationMixin):
def __init__(self, model_dir):
hf_config = PretrainedConfig.from_pretrained(model_dir)
+ # Ensure text_config is a proper object, not a dict (fixes qwen3_5_moe + transformers compat)
+ if isinstance(getattr(hf_config, 'text_config', None), dict):
+ hf_config.text_config = PretrainedConfig(**hf_config.text_config)
super().__init__(hf_config)
exl3_config = Config.from_directory(model_dir)
@@ -201,26 +202,11 @@ class Exllamav3HF(PreTrainedModel, GenerationMixin):
}
).to(input_ids.device).float()
else:
- # When processing with labels, handle as a complete sequence
- # Process in chunks if the number of tokens is large
- tokens_to_process = seq_tensor
- all_logits = None
-
- for i in range(0, tokens_to_process.shape[0], max_chunk_size):
- chunk = tokens_to_process[i:i + max_chunk_size]
- chunk_logits = self.ex_model.forward(
- input_ids=chunk.view(1, -1),
- params={
- "attn_mode": "flash_attn_nc",
- }
- ).float()
-
- if all_logits is None:
- all_logits = chunk_logits
- else:
- all_logits = torch.cat([all_logits, chunk_logits], dim=1)
-
- logits = all_logits
+ # Labels path: single pass without cache for correct logits
+ logits = self.ex_model.forward(
+ input_ids=seq_tensor.view(1, -1),
+ params={"attn_mode": "flash_attn_nc"}
+ ).float().cpu()
if is_negative:
self.past_seq_negative = seq_tensor
diff --git a/modules/extensions.py b/modules/extensions.py
index dd327882..afe847f0 100644
--- a/modules/extensions.py
+++ b/modules/extensions.py
@@ -1,13 +1,10 @@
import importlib
import importlib.util
import sys
-import traceback
from functools import partial
from inspect import signature
from pathlib import Path
-import gradio as gr
-
import modules.shared as shared
from modules.logging_colors import logger
@@ -35,8 +32,7 @@ def load_extensions():
if name not in available_extensions:
continue
- if name != 'api':
- logger.info(f'Loading the extension "{name}"')
+ logger.info(f'Loading the extension "{name}"')
try:
# Prefer user extension, fall back to system extension
@@ -77,8 +73,7 @@ def load_extensions():
raise
except Exception:
- logger.error(f'Failed to load the extension "{name}".')
- traceback.print_exc()
+ logger.exception(f'Failed to load the extension "{name}".')
# This iterator returns the extensions in the order specified in the command-line
@@ -196,24 +191,23 @@ def _apply_custom_generate_reply():
def _apply_custom_css():
- all_css = ''
- for extension, _ in iterator():
- if hasattr(extension, 'custom_css'):
- all_css += getattr(extension, 'custom_css')()
-
- return all_css
+ return ''.join(
+ getattr(extension, 'custom_css')()
+ for extension, _ in iterator()
+ if hasattr(extension, 'custom_css')
+ )
def _apply_custom_js():
- all_js = ''
- for extension, _ in iterator():
- if hasattr(extension, 'custom_js'):
- all_js += getattr(extension, 'custom_js')()
-
- return all_js
+ return ''.join(
+ getattr(extension, 'custom_js')()
+ for extension, _ in iterator()
+ if hasattr(extension, 'custom_js')
+ )
def create_extensions_block():
+ import gradio as gr
to_display = []
for extension, name in iterator():
if hasattr(extension, "ui") and not (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
@@ -228,6 +222,7 @@ def create_extensions_block():
def create_extensions_tabs():
+ import gradio as gr
for extension, name in iterator():
if hasattr(extension, "ui") and (hasattr(extension, 'params') and extension.params.get('is_tab', False)):
display_name = getattr(extension, 'params', {}).get('display_name', name)
diff --git a/modules/grammar/__init__.py b/modules/grammar/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/modules/html_generator.py b/modules/html_generator.py
index 472a9ea0..8f3f261f 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -10,6 +10,7 @@ import markdown
from PIL import Image, ImageOps
from modules import shared
+from modules.reasoning import extract_reasoning
from modules.sane_markdown_lists import SaneListExtension
from modules.utils import get_available_chat_styles
@@ -108,69 +109,41 @@ def replace_blockquote(m):
return m.group().replace('\n', '\n> ').replace('\\begin{blockquote}', '').replace('\\end{blockquote}', '')
-# Thinking block format definitions: (start_tag, end_tag, content_start_tag)
-# Use None for start_tag to match from beginning (end-only formats should be listed last)
-THINKING_FORMATS = [
- ('', '', None),
- ('<|channel|>analysis<|message|>', '<|end|>', '<|start|>assistant<|channel|>final<|message|>'),
- ('', '', None),
- ('<|think|>', '<|end|>', '<|content|>'), # Solar Open
- ('Thinking Process:', '', None), # Qwen3.5 verbose thinking outside tags
- (None, '', None), # End-only variant (e.g., Qwen3-next)
-]
-
-
def extract_thinking_block(string):
- """Extract thinking blocks from the beginning of a string."""
- if not string:
- return None, string
-
- for start_tag, end_tag, content_tag in THINKING_FORMATS:
- end_esc = html.escape(end_tag)
- content_esc = html.escape(content_tag) if content_tag else None
-
- if start_tag is None:
- # End-only format: require end tag, start from beginning
- end_pos = string.find(end_esc)
- if end_pos == -1:
- continue
- thought_start = 0
- else:
- # Normal format: require start tag
- start_esc = html.escape(start_tag)
- start_pos = string.find(start_esc)
- if start_pos == -1:
- continue
- thought_start = start_pos + len(start_esc)
- end_pos = string.find(end_esc, thought_start)
-
- if end_pos == -1:
- # End tag missing - check if content tag can serve as fallback
- if content_esc:
- content_pos = string.find(content_esc, thought_start)
- if content_pos != -1:
- thought_end = content_pos
- content_start = content_pos + len(content_esc)
- else:
- thought_end = len(string)
- content_start = len(string)
- else:
- thought_end = len(string)
- content_start = len(string)
- else:
- thought_end = end_pos
- if content_esc:
- content_pos = string.find(content_esc, end_pos)
- content_start = content_pos + len(content_esc) if content_pos != -1 else end_pos + len(end_esc)
- else:
- content_start = end_pos + len(end_esc)
-
- return string[thought_start:thought_end], string[content_start:]
-
- return None, string
+ """Extract thinking blocks from the beginning of an HTML-escaped string."""
+ return extract_reasoning(string, html_escaped=True)
-def build_thinking_block(thinking_content, message_id, has_remaining_content):
+
+def build_tool_call_block(header, body, message_id, index):
+ """Build HTML for a tool call accordion block."""
+ block_id = f"tool-call-{message_id}-{index}"
+
+ if body == '...':
+ # Pending placeholder — no expandable body, just title with ellipsis
+ return f'''
+
+
+
+ '''
+
+ # Build a plain directly to avoid highlight.js auto-detection
+ escaped_body = html.escape(body)
+ return f'''
+
+
+
+
+ '''
+
+
+def build_thinking_block(thinking_content, message_id, has_remaining_content, thinking_index=0):
"""Build HTML for a thinking block."""
if thinking_content is None:
return None
@@ -179,7 +152,7 @@ def build_thinking_block(thinking_content, message_id, has_remaining_content):
thinking_html = process_markdown_content(thinking_content)
# Generate unique ID for the thinking block
- block_id = f"thinking-{message_id}-0"
+ block_id = f"thinking-{message_id}-{thinking_index}"
# Check if thinking is complete or still in progress
is_streaming = not has_remaining_content
@@ -344,6 +317,9 @@ def process_markdown_content(string):
# Unescape backslashes
html_output = html_output.replace('\\\\', '\\')
+ # Wrap tables in a scrollable div
+ html_output = html_output.replace('
')
+
return html_output
@@ -360,24 +336,66 @@ def convert_to_markdown(string, message_id=None):
if message_id is None:
message_id = "unknown"
- # Extract different components from the string
- thinking_content, remaining_content = extract_thinking_block(string)
+ # Find tool call blocks by position, then process the text segments
+ # between them using extract_thinking_block (which supports all
+ # THINKING_FORMATS, including end-only variants like Qwen's).
+ tool_call_pattern = re.compile(r'(.*?)\n(.*?)\n', re.DOTALL)
+ tool_calls = list(tool_call_pattern.finditer(string))
- # Build individual HTML blocks
- blocks = []
+ if not tool_calls:
+ # No tool calls — use original single-pass extraction
+ thinking_content, remaining_content = extract_thinking_block(string)
+ blocks = []
+ thinking_html = build_thinking_block(thinking_content, message_id, bool(remaining_content))
+ if thinking_html:
+ blocks.append(thinking_html)
- # Add thinking block if present
- thinking_html = build_thinking_block(thinking_content, message_id, bool(remaining_content))
- if thinking_html:
- blocks.append(thinking_html)
+ main_html = build_main_content_block(remaining_content)
+ if main_html:
+ blocks.append(main_html)
- # Add main content block
- main_html = build_main_content_block(remaining_content)
- if main_html:
- blocks.append(main_html)
+ return ''.join(blocks)
- # Assemble all blocks into final HTML
- return ''.join(blocks)
+ # Split string into text segments around tool_call blocks and
+ # run extract_thinking_block on each segment for full format support.
+ html_parts = []
+ last_end = 0
+ tool_idx = 0
+ think_idx = 0
+
+ def process_text_segment(text, is_last_segment):
+ """Process a text segment between tool_call blocks for thinking content."""
+ nonlocal think_idx
+ if not text.strip():
+ return
+
+ while text.strip():
+ thinking_content, remaining = extract_thinking_block(text)
+ if thinking_content is None:
+ break
+ has_remaining = bool(remaining.strip()) or not is_last_segment
+ html_parts.append(build_thinking_block(thinking_content, message_id, has_remaining, think_idx))
+ think_idx += 1
+ text = remaining
+
+ if text.strip():
+ html_parts.append(process_markdown_content(text))
+
+ for tc in tool_calls:
+ # Process text before this tool_call
+ process_text_segment(string[last_end:tc.start()], is_last_segment=False)
+
+ # Add tool call accordion
+ header = tc.group(1).strip()
+ body = tc.group(2).strip()
+ html_parts.append(build_tool_call_block(header, body, message_id, tool_idx))
+ tool_idx += 1
+ last_end = tc.end()
+
+ # Process text after the last tool_call
+ process_text_segment(string[last_end:], is_last_segment=True)
+
+ return ''.join(html_parts)
def convert_to_markdown_wrapped(string, message_id=None, use_cache=True):
@@ -435,6 +453,7 @@ branch_svg = ''''''
info_svg = ''''''
info_svg_small = ''''''
+tool_svg_small = ''''''
attachment_svg = ''''''
copy_button = f''
diff --git a/modules/image_utils.py b/modules/image_utils.py
index d2809fef..b3138790 100644
--- a/modules/image_utils.py
+++ b/modules/image_utils.py
@@ -77,7 +77,18 @@ def process_message_content(content: Any) -> Tuple[str, List[Image.Image]]:
# Support external URLs
try:
import requests
- response = requests.get(image_url, timeout=10)
+ from urllib.parse import urljoin
+ from modules.web_search import _validate_url
+ _validate_url(image_url)
+ url = image_url
+ for _ in range(5):
+ response = requests.get(url, timeout=10, allow_redirects=False)
+ if response.is_redirect and 'Location' in response.headers:
+ url = urljoin(url, response.headers['Location'])
+ _validate_url(url)
+ else:
+ break
+
response.raise_for_status()
image_data = response.content
image = Image.open(io.BytesIO(image_data))
diff --git a/modules/llama_cpp_server.py b/modules/llama_cpp_server.py
index 12ff173e..34080466 100644
--- a/modules/llama_cpp_server.py
+++ b/modules/llama_cpp_server.py
@@ -1,6 +1,7 @@
import json
import os
import pprint
+import shlex
import re
import socket
import subprocess
@@ -10,7 +11,6 @@ import time
from pathlib import Path
from typing import Any, List
-import llama_cpp_binaries
import requests
from modules import shared
@@ -36,6 +36,7 @@ class LlamaServer:
self.process = None
self.session = requests.Session()
self.vocabulary_size = None
+ self.n_ctx = None
self.bos_token = ""
self.last_prompt_token_count = 0
@@ -129,13 +130,24 @@ class LlamaServer:
# places it at the end of the chain regardless of position, so we
# activate it based on the parameter value rather than sampler order.
if state.get("adaptive_target", 0) > 0:
- filtered_samplers.append("adaptive-p")
+ filtered_samplers.append("adaptive_p")
payload["samplers"] = filtered_samplers
+ logit_bias = []
if state['custom_token_bans']:
- to_ban = [[int(token_id), False] for token_id in state['custom_token_bans'].split(',')]
- payload["logit_bias"] = to_ban
+ logit_bias.extend([[int(token_id.strip()), False] for token_id in state['custom_token_bans'].split(',') if token_id.strip()])
+
+ if state.get('logit_bias'):
+ for token_id_str, bias in state['logit_bias'].items():
+ logit_bias.append([int(token_id_str), bias])
+
+ if logit_bias:
+ payload["logit_bias"] = logit_bias
+
+ n_probs = state.get('logprobs', 0)
+ if n_probs and n_probs > 0:
+ payload["n_probs"] = n_probs
return payload
@@ -215,6 +227,7 @@ class LlamaServer:
response.raise_for_status() # Raise an exception for HTTP errors
full_text = ""
+ self.last_completion_probabilities = []
# Process the streaming response
stop_event = state.get('stop_event')
@@ -240,6 +253,10 @@ class LlamaServer:
full_text += data['content']
yield full_text
+ # Capture logprobs if present
+ if 'completion_probabilities' in data:
+ self.last_completion_probabilities.extend(data['completion_probabilities'])
+
# Check if generation is complete
if data.get('stop', False):
break
@@ -293,8 +310,45 @@ class LlamaServer:
else:
raise Exception(f"Unexpected response format: 'completion_probabilities' not found in {result}")
+ def get_prompt_logprob_entries(self, token_ids, n_probs=5, prompt=""):
+ """Get logprob entries for prompt tokens via a single n_predict=0 request.
+
+ Requires llama.cpp server with prompt_logprobs support.
+ Returns entries in the standard format for format_completion_logprobs().
+ """
+ token_ids_list = token_ids.tolist() if hasattr(token_ids, 'tolist') else list(token_ids)
+
+ url = f"http://127.0.0.1:{self.port}/completion"
+ payload = {
+ "prompt": token_ids_list,
+ "n_predict": 0,
+ "n_probs": n_probs,
+ "prompt_logprobs": True,
+ "stream": False,
+ "cache_prompt": False,
+ }
+
+ response = self.session.post(url, json=payload)
+ result = response.json()
+
+ prompt_probs = result.get("prompt_probabilities", [])
+ if not prompt_probs:
+ return []
+
+ # Null first token (no conditioning context); use empty string for BOS
+ # or tokens that don't appear at the start of the prompt text.
+ first_token_str = self.decode([token_ids_list[0]])
+ if self.bos_token and first_token_str == self.bos_token:
+ first_token_str = ""
+ elif not prompt.startswith(first_token_str):
+ first_token_str = ""
+
+ entries = [{"token": first_token_str, "null_logprob": True}]
+ entries.extend(prompt_probs)
+ return entries
+
def _get_vocabulary_size(self):
- """Get and store the model's maximum context length."""
+ """Get and store the model's vocabulary size."""
url = f"http://127.0.0.1:{self.port}/v1/models"
response = self.session.get(url).json()
@@ -304,12 +358,17 @@ class LlamaServer:
self.vocabulary_size = model_info["meta"]["n_vocab"]
def _get_bos_token(self):
- """Get and store the model's BOS token."""
+ """Get and store the model's BOS token and context size."""
url = f"http://127.0.0.1:{self.port}/props"
response = self.session.get(url).json()
if "bos_token" in response:
self.bos_token = response["bos_token"]
+ # Get actual n_ctx from the server (important when --fit auto-selects it)
+ n_ctx = response.get("default_generation_settings", {}).get("n_ctx")
+ if n_ctx:
+ self.n_ctx = n_ctx
+
def _is_port_available(self, port):
"""Check if a port is available for use."""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -334,7 +393,16 @@ class LlamaServer:
"""Start the llama.cpp server and wait until it's ready."""
# Determine the server path
if self.server_path is None:
- self.server_path = llama_cpp_binaries.get_binary_path()
+ if shared.args.ik:
+ try:
+ import ik_llama_cpp_binaries
+ except ImportError:
+ raise ImportError("--ik requires the ik_llama_cpp_binaries package. Install it with: pip install ")
+
+ self.server_path = ik_llama_cpp_binaries.get_binary_path()
+ else:
+ import llama_cpp_binaries
+ self.server_path = llama_cpp_binaries.get_binary_path()
# Build the command
cmd = [
@@ -349,11 +417,14 @@ class LlamaServer:
if shared.args.ctx_size > 0:
cmd += ["--ctx-size", str(shared.args.ctx_size)]
+ elif shared.args.gpu_layers >= 0:
+ cmd += ["--ctx-size", "8192"]
if shared.args.gpu_layers >= 0:
cmd += ["--gpu-layers", str(shared.args.gpu_layers), "--fit", "off"]
else:
cmd += ["--fit", "on"]
+ cmd += ["--fit-ctx", "8192"]
if shared.args.fit_target:
cmd += ["--fit-target", shared.args.fit_target]
@@ -379,10 +450,6 @@ class LlamaServer:
if shared.args.cache_type != "fp16" and shared.args.cache_type in llamacpp_valid_cache_types:
cmd += ["--cache-type-k", shared.args.cache_type, "--cache-type-v", shared.args.cache_type]
cache_type = shared.args.cache_type
- if shared.args.compress_pos_emb != 1:
- cmd += ["--rope-freq-scale", str(1.0 / shared.args.compress_pos_emb)]
- if shared.args.rope_freq_base > 0:
- cmd += ["--rope-freq-base", str(shared.args.rope_freq_base)]
if shared.args.mmproj not in [None, 'None']:
path = Path(shared.args.mmproj)
if not path.exists():
@@ -425,21 +492,32 @@ class LlamaServer:
elif extra_flags.startswith("'") and extra_flags.endswith("'"):
extra_flags = extra_flags[1:-1].strip()
- for flag_item in extra_flags.split(','):
- flag_item = flag_item.strip()
- if '=' in flag_item:
- flag, value = flag_item.split('=', 1)
- flag = flag.strip()
- value = value.strip()
- if len(flag) <= 3:
- cmd += [f"-{flag}", value]
+ if extra_flags.startswith('-'):
+ # New literal format: "--jinja --rpc 1222,1222"
+ cmd += shlex.split(extra_flags)
+ else:
+ # Legacy format: "flag1=value1,flag2,flag3=value3"
+ long_form_only = {'rpc', 'fit', 'pos', 'ppl'}
+
+ for flag_item in extra_flags.split(','):
+ flag_item = flag_item.strip()
+ if '=' in flag_item:
+ flag, value = flag_item.split('=', 1)
+ flag = flag.strip()
+ value = value.strip()
+ if len(flag) <= 3 and flag not in long_form_only:
+ cmd += [f"-{flag}", value]
+ else:
+ cmd += [f"--{flag}", value]
else:
- cmd += [f"--{flag}", value]
- else:
- if len(flag_item) <= 3:
- cmd.append(f"-{flag_item}")
- else:
- cmd.append(f"--{flag_item}")
+ if len(flag_item) <= 3 and flag_item not in long_form_only:
+ cmd.append(f"-{flag_item}")
+ else:
+ cmd.append(f"--{flag_item}")
+
+ # Patch flags for ik_llama.cpp compatibility
+ if shared.args.ik:
+ cmd = _patch_cmd_for_ik(cmd)
env = os.environ.copy()
if os.name == 'posix':
@@ -455,7 +533,7 @@ class LlamaServer:
print()
gpu_layers_str = "auto" if shared.args.gpu_layers < 0 else str(shared.args.gpu_layers)
- ctx_size_str = "auto" if shared.args.ctx_size == 0 else str(shared.args.ctx_size)
+ ctx_size_str = "auto" if shared.args.ctx_size == 0 and shared.args.gpu_layers < 0 else str(shared.args.ctx_size or 8192)
logger.info(f"Using gpu_layers={gpu_layers_str} | ctx_size={ctx_size_str} | cache_type={cache_type}")
# Start the server with pipes for output
self.process = subprocess.Popen(
@@ -471,9 +549,8 @@ class LlamaServer:
health_url = f"http://127.0.0.1:{self.port}/health"
while True:
# Check if process is still alive
- if self.process.poll() is not None:
- # Process has terminated
- exit_code = self.process.poll()
+ exit_code = self.process.poll()
+ if exit_code is not None:
raise RuntimeError(f"Server process terminated unexpectedly with exit code: {exit_code}")
try:
@@ -579,3 +656,49 @@ def filter_stderr_with_progress(process_stderr):
process_stderr.close()
except Exception:
pass
+
+
+def _patch_cmd_for_ik(cmd):
+ """
+ Rewrite upstream llama.cpp flags to ik_llama.cpp equivalents:
+ --no-webui → --webui none
+ --fit off → (removed)
+ --fit on / --fit-ctx → --fit (bare flag)
+ --fit-target → --fit-margin
+ --cache-reuse → (removed, unsupported)
+ --swa-full → (removed, unsupported)
+ """
+ # Add Hadamard KV cache rotation when using quantized cache types.
+ # This significantly improves quantized cache quality (especially q4_0)
+ # and is a no-op for MLA models like DeepSeek.
+ if shared.args.cache_type in ("q8_0", "q4_0"):
+ cmd += ["-khad", "-vhad"]
+
+ patched = []
+ i = 0
+ while i < len(cmd):
+ arg = cmd[i]
+
+ if arg == "--no-webui":
+ patched += ["--webui", "none"]
+ elif arg == "--fit" and i + 1 < len(cmd) and cmd[i + 1] in ("on", "off"):
+ val = cmd[i + 1]
+ i += 1
+ if val == "on":
+ patched.append("--fit")
+ # "off" → drop entirely
+ elif arg == "--fit-ctx":
+ patched.append("--fit")
+ i += 1 # skip the value
+ elif arg == "--fit-target":
+ patched.append("--fit-margin")
+ elif arg == "--cache-reuse":
+ i += 1 # skip the value
+ elif arg == "--swa-full":
+ pass # bare flag, just drop it
+ else:
+ patched.append(arg)
+
+ i += 1
+
+ return patched
diff --git a/modules/loaders.py b/modules/loaders.py
index 42a5ff1c..31b1b51a 100644
--- a/modules/loaders.py
+++ b/modules/loaders.py
@@ -1,8 +1,6 @@
import functools
from collections import OrderedDict
-import gradio as gr
-
loaders_and_params = OrderedDict({
'llama.cpp': [
'gpu_layers',
@@ -17,13 +15,12 @@ loaders_and_params = OrderedDict({
'tensor_split',
'extra_flags',
'streaming_llm',
- 'rope_freq_base',
- 'compress_pos_emb',
'row_split',
'no_kv_offload',
'no_mmap',
'mlock',
'numa',
+ 'ik',
'parallel',
'model_draft',
'draft_max',
@@ -43,8 +40,6 @@ loaders_and_params = OrderedDict({
'Transformers': [
'gpu_split',
'cpu_memory',
- 'alpha_value',
- 'compress_pos_emb',
'compute_dtype',
'quant_type',
'load_in_8bit',
@@ -71,7 +66,6 @@ loaders_and_params = OrderedDict({
'gpu_split',
'model_draft',
'draft_max',
- 'ctx_size_draft',
'speculative_decoding_accordion',
'enable_tp',
'tp_backend',
@@ -208,6 +202,7 @@ loaders_samplers = {
'ban_eos_token',
'add_bos_token',
'enable_thinking',
+ 'reasoning_effort',
'seed',
'skip_special_tokens',
},
@@ -244,6 +239,7 @@ loaders_samplers = {
'reasoning_effort',
'seed',
'sampler_priority',
+ 'custom_token_bans',
'dry_sequence_breakers',
'grammar_string',
'grammar_file_row',
@@ -277,6 +273,7 @@ def list_all_samplers():
def blacklist_samplers(loader, dynamic_temperature):
+ import gradio as gr
all_samplers = list_all_samplers()
output = []
@@ -294,15 +291,77 @@ def blacklist_samplers(loader, dynamic_temperature):
@functools.cache
def get_all_params():
+ from modules import shared
all_params = set()
for k in loaders_and_params:
for el in loaders_and_params[k]:
all_params.add(el)
+ if shared.args.portable:
+ all_params.discard('ik')
+
return sorted(all_params)
+@functools.cache
+def list_model_elements():
+ elements = [
+ 'filter_by_loader',
+ 'loader',
+ 'cpu_memory',
+ 'gpu_layers',
+ 'fit_target',
+ 'cpu_moe',
+ 'threads',
+ 'threads_batch',
+ 'batch_size',
+ 'ubatch_size',
+ 'ctx_size',
+ 'cache_type',
+ 'tensor_split',
+ 'extra_flags',
+ 'streaming_llm',
+ 'gpu_split',
+ 'compute_dtype',
+ 'quant_type',
+ 'load_in_8bit',
+ 'load_in_4bit',
+ 'attn_implementation',
+ 'cpu',
+ 'disk',
+ 'row_split',
+ 'no_kv_offload',
+ 'no_mmap',
+ 'mlock',
+ 'numa',
+ 'parallel',
+ 'use_double_quant',
+ 'bf16',
+ 'enable_tp',
+ 'tp_backend',
+ 'cfg_cache',
+ 'no_use_fast',
+ 'model_draft',
+ 'draft_max',
+ 'gpu_layers_draft',
+ 'device_draft',
+ 'ctx_size_draft',
+ 'spec_type',
+ 'spec_ngram_size_n',
+ 'spec_ngram_size_m',
+ 'spec_ngram_min_hits',
+ 'mmproj',
+ ]
+
+ from modules import shared
+ if not shared.args.portable:
+ elements.append('ik')
+
+ return elements
+
+
def make_loader_params_visible(loader):
+ import gradio as gr
params = []
all_params = get_all_params()
if loader in loaders_and_params:
diff --git a/modules/logits.py b/modules/logits.py
index 2d066c09..473f5890 100644
--- a/modules/logits.py
+++ b/modules/logits.py
@@ -1,11 +1,9 @@
import time
-import traceback
import numpy as np
from modules import models, shared
from modules.logging_colors import logger
-from modules.models import load_model
from modules.text_generation import generate_reply
from modules.utils import check_model_loaded
@@ -13,8 +11,7 @@ global_scores = None
def get_next_logits(*args, **kwargs):
- if shared.args.idle_timeout > 0 and shared.model is None and shared.model_name not in [None, 'None']:
- shared.model, shared.tokenizer = load_model(shared.model_name)
+ models.load_model_if_idle_unloaded()
needs_lock = not args[2] # use_samplers
if needs_lock:
@@ -23,7 +20,7 @@ def get_next_logits(*args, **kwargs):
try:
result = _get_next_logits(*args, **kwargs)
except Exception:
- traceback.print_exc()
+ logger.exception("Failed to get next logits")
result = None
if needs_lock:
diff --git a/modules/models.py b/modules/models.py
index 48d68b0b..e997d2d8 100644
--- a/modules/models.py
+++ b/modules/models.py
@@ -1,4 +1,5 @@
import sys
+import threading
import time
import modules.shared as shared
@@ -7,6 +8,15 @@ from modules.models_settings import get_model_metadata
from modules.utils import resolve_model_path
last_generation_time = time.time()
+active_generation_count = 0
+_generation_count_lock = threading.Lock()
+
+
+def load_model_if_idle_unloaded():
+ global last_generation_time
+ if shared.args.idle_timeout > 0 and shared.model is None and shared.model_name not in [None, 'None']:
+ shared.model, shared.tokenizer = load_model(shared.model_name)
+ last_generation_time = time.time()
def load_model(model_name, loader=None):
@@ -38,6 +48,9 @@ def load_model(model_name, loader=None):
sampler_hijack.hijack_samplers()
shared.args.loader = loader
+ if loader != 'llama.cpp' and shared.args.ctx_size == 0:
+ shared.args.ctx_size = 8192
+
output = load_func_map[loader](model_name)
if type(output) is tuple:
model, tokenizer = output
@@ -54,6 +67,8 @@ def load_model(model_name, loader=None):
if loader.lower().startswith('exllama') or loader.lower().startswith('tensorrt') or loader == 'llama.cpp':
if shared.args.ctx_size > 0:
shared.settings['truncation_length'] = shared.args.ctx_size
+ elif loader == 'llama.cpp' and hasattr(model, 'n_ctx') and model.n_ctx:
+ shared.settings['truncation_length'] = model.n_ctx
shared.is_multimodal = False
if loader.lower() in ('exllamav3', 'llama.cpp') and hasattr(model, 'is_multimodal'):
@@ -61,8 +76,7 @@ def load_model(model_name, loader=None):
logger.info(f"Loaded \"{model_name}\" in {(time.time()-t0):.2f} seconds.")
logger.info(f"LOADER: \"{loader}\"")
- logger.info(f"TRUNCATION LENGTH: {shared.settings['truncation_length']}")
- logger.info(f"INSTRUCTION TEMPLATE: \"{metadata['instruction_template']}\"")
+ logger.info(f"CONTEXT LENGTH: {shared.settings['truncation_length']}")
return model, tokenizer
@@ -154,7 +168,10 @@ def unload_model_if_idle():
while True:
shared.generation_lock.acquire()
try:
- if time.time() - last_generation_time > shared.args.idle_timeout * 60:
+ with _generation_count_lock:
+ is_active = active_generation_count > 0
+
+ if not is_active and time.time() - last_generation_time > shared.args.idle_timeout * 60:
if shared.model is not None:
logger.info("Unloading the model for inactivity.")
unload_model(keep_model_name=True)
diff --git a/modules/models_settings.py b/modules/models_settings.py
index 472871ce..eafa0581 100644
--- a/modules/models_settings.py
+++ b/modules/models_settings.py
@@ -4,10 +4,9 @@ import re
from math import floor
from pathlib import Path
-import gradio as gr
import yaml
-from modules import chat, loaders, metadata_gguf, shared, ui
+from modules import loaders, metadata_gguf, shared
from modules.logging_colors import logger
from modules.utils import resolve_model_path
@@ -16,9 +15,6 @@ def get_fallback_settings():
return {
'bf16': False,
'ctx_size': 8192,
- 'rope_freq_base': 0,
- 'compress_pos_emb': 1,
- 'alpha_value': 1,
'truncation_length': shared.settings['truncation_length'],
'truncation_length_info': shared.settings['truncation_length'],
'skip_special_tokens': shared.settings['skip_special_tokens'],
@@ -27,18 +23,14 @@ def get_fallback_settings():
def get_model_metadata(model):
model_path = resolve_model_path(model)
- model_settings = {}
- # Get settings from user_data/models/config.yaml and user_data/models/config-user.yaml
- settings = shared.model_config
- for pat in settings:
- if re.match(pat.lower(), Path(model).name.lower()):
- for k in settings[pat]:
- model_settings[k] = settings[pat][k]
+ # Fallback settings
+ model_settings = get_fallback_settings()
path = model_path / 'config.json'
if path.exists():
- hf_metadata = json.loads(open(path, 'r', encoding='utf-8').read())
+ with open(path, 'r', encoding='utf-8') as f:
+ hf_metadata = json.loads(f.read())
else:
hf_metadata = None
@@ -68,14 +60,8 @@ def get_model_metadata(model):
for k in metadata:
if k.endswith('.context_length'):
- model_settings['ctx_size'] = min(metadata[k], 8192)
+ model_settings['ctx_size'] = 0
model_settings['truncation_length_info'] = metadata[k]
- elif k.endswith('rope.freq_base'):
- model_settings['rope_freq_base'] = metadata[k]
- elif k.endswith('rope.scale_linear'):
- model_settings['compress_pos_emb'] = metadata[k]
- elif k.endswith('rope.scaling.factor'):
- model_settings['compress_pos_emb'] = metadata[k]
elif k.endswith('.block_count'):
model_settings['gpu_layers'] = -1
model_settings['max_gpu_layers'] = metadata[k] + 1
@@ -103,7 +89,7 @@ def get_model_metadata(model):
else:
# Transformers metadata
if hf_metadata is not None:
- metadata = json.loads(open(path, 'r', encoding='utf-8').read())
+ metadata = hf_metadata
if 'pretrained_config' in metadata:
metadata = metadata['pretrained_config']
@@ -120,15 +106,6 @@ def get_model_metadata(model):
model_settings['ctx_size'] = min(value, 8192)
break
- if 'rope_theta' in metadata:
- model_settings['rope_freq_base'] = metadata['rope_theta']
- elif 'attn_config' in metadata and 'rope_theta' in metadata['attn_config']:
- model_settings['rope_freq_base'] = metadata['attn_config']['rope_theta']
-
- if 'rope_scaling' in metadata and isinstance(metadata['rope_scaling'], dict) and all(key in metadata['rope_scaling'] for key in ('type', 'factor')):
- if metadata['rope_scaling']['type'] == 'linear':
- model_settings['compress_pos_emb'] = metadata['rope_scaling']['factor']
-
if 'torch_dtype' in metadata and metadata['torch_dtype'] == 'bfloat16':
model_settings['bf16'] = True
@@ -153,7 +130,8 @@ def get_model_metadata(model):
# 3. Fall back to tokenizer_config.json metadata
if path.exists():
- metadata = json.loads(open(path, 'r', encoding='utf-8').read())
+ with open(path, 'r', encoding='utf-8') as f:
+ metadata = json.loads(f.read())
# Only read from metadata if we haven't already loaded from .jinja or .json
if template is None and 'chat_template' in metadata:
@@ -182,10 +160,6 @@ def get_model_metadata(model):
if 'instruction_template' not in model_settings:
model_settings['instruction_template'] = 'Alpaca'
- # Ignore rope_freq_base if set to the default value
- if 'rope_freq_base' in model_settings and model_settings['rope_freq_base'] == 10000:
- model_settings.pop('rope_freq_base')
-
# Apply user settings from user_data/models/config-user.yaml
settings = shared.user_config
for pat in settings:
@@ -199,7 +173,7 @@ def get_model_metadata(model):
# Load instruction template if defined by name rather than by value
if model_settings['instruction_template'] != 'Custom (obtained from model metadata)':
- model_settings['instruction_template_str'] = chat.load_instruction_template(model_settings['instruction_template'])
+ model_settings['instruction_template_str'] = load_instruction_template(model_settings['instruction_template'])
return model_settings
@@ -228,7 +202,7 @@ def update_model_parameters(state, initial=False):
'''
UI: update the command-line arguments based on the interface values
'''
- elements = ui.list_model_elements() # the names of the parameters
+ elements = loaders.list_model_elements() # the names of the parameters
for i, element in enumerate(elements):
if element not in state:
@@ -248,6 +222,7 @@ def apply_model_settings_to_state(model, state):
'''
UI: update the state variable with the model settings
'''
+ import gradio as gr
model_settings = get_model_metadata(model)
if 'loader' in model_settings:
loader = model_settings.pop('loader')
@@ -290,7 +265,7 @@ def save_model_settings(model, state):
if model_regex not in user_config:
user_config[model_regex] = {}
- for k in ui.list_model_elements():
+ for k in loaders.list_model_elements():
if k == 'loader' or k in loaders.loaders_and_params[state['loader']]:
user_config[model_regex][k] = state[k]
@@ -419,3 +394,103 @@ def update_gpu_layers_and_vram(loader, model, gpu_layers, ctx_size, cache_type):
vram_usage = estimate_vram(model, gpu_layers, ctx_size, cache_type)
return f"Estimated VRAM to load the model: {vram_usage:.0f} MiB
"
+
+
+def load_instruction_template(template):
+ if template == 'None':
+ return ''
+
+ for filepath in [shared.user_data_dir / 'instruction-templates' / f'{template}.yaml', shared.user_data_dir / 'instruction-templates' / 'Alpaca.yaml']:
+ if filepath.exists():
+ break
+ else:
+ return ''
+
+ with open(filepath, 'r', encoding='utf-8') as f:
+ file_contents = f.read()
+ data = yaml.safe_load(file_contents)
+ if 'instruction_template' in data:
+ return data['instruction_template']
+ else:
+ return _jinja_template_from_old_format(data)
+
+
+def _jinja_template_from_old_format(params, verbose=False):
+ MASTER_TEMPLATE = """
+{%- set ns = namespace(found=false) -%}
+{%- for message in messages -%}
+ {%- if message['role'] == 'system' -%}
+ {%- set ns.found = true -%}
+ {%- endif -%}
+{%- endfor -%}
+{%- if not ns.found -%}
+ {{- '<|PRE-SYSTEM|>' + '<|SYSTEM-MESSAGE|>' + '<|POST-SYSTEM|>' -}}
+{%- endif %}
+{%- for message in messages %}
+ {%- if message['role'] == 'system' -%}
+ {{- '<|PRE-SYSTEM|>' + message['content'] + '<|POST-SYSTEM|>' -}}
+ {%- else -%}
+ {%- if message['role'] == 'user' -%}
+ {{-'<|PRE-USER|>' + message['content'] + '<|POST-USER|>'-}}
+ {%- else -%}
+ {{-'<|PRE-ASSISTANT|>' + message['content'] + '<|POST-ASSISTANT|>' -}}
+ {%- endif -%}
+ {%- endif -%}
+{%- endfor -%}
+{%- if add_generation_prompt -%}
+ {{-'<|PRE-ASSISTANT-GENERATE|>'-}}
+{%- endif -%}
+"""
+
+ if 'context' in params and '<|system-message|>' in params['context']:
+ pre_system = params['context'].split('<|system-message|>')[0]
+ post_system = params['context'].split('<|system-message|>')[1]
+ else:
+ pre_system = ''
+ post_system = ''
+
+ pre_user = params['turn_template'].split('<|user-message|>')[0].replace('<|user|>', params['user'])
+ post_user = params['turn_template'].split('<|user-message|>')[1].split('<|bot|>')[0]
+
+ pre_assistant = '<|bot|>' + params['turn_template'].split('<|bot-message|>')[0].split('<|bot|>')[1]
+ pre_assistant = pre_assistant.replace('<|bot|>', params['bot'])
+ post_assistant = params['turn_template'].split('<|bot-message|>')[1]
+
+ def preprocess(string):
+ return string.replace('\n', '\\n').replace('\'', '\\\'')
+
+ pre_system = preprocess(pre_system)
+ post_system = preprocess(post_system)
+ pre_user = preprocess(pre_user)
+ post_user = preprocess(post_user)
+ pre_assistant = preprocess(pre_assistant)
+ post_assistant = preprocess(post_assistant)
+
+ if verbose:
+ print(
+ '\n',
+ repr(pre_system) + '\n',
+ repr(post_system) + '\n',
+ repr(pre_user) + '\n',
+ repr(post_user) + '\n',
+ repr(pre_assistant) + '\n',
+ repr(post_assistant) + '\n',
+ )
+
+ result = MASTER_TEMPLATE
+ if 'system_message' in params:
+ result = result.replace('<|SYSTEM-MESSAGE|>', preprocess(params['system_message']))
+ else:
+ result = result.replace('<|SYSTEM-MESSAGE|>', '')
+
+ result = result.replace('<|PRE-SYSTEM|>', pre_system)
+ result = result.replace('<|POST-SYSTEM|>', post_system)
+ result = result.replace('<|PRE-USER|>', pre_user)
+ result = result.replace('<|POST-USER|>', post_user)
+ result = result.replace('<|PRE-ASSISTANT|>', pre_assistant)
+ result = result.replace('<|PRE-ASSISTANT-GENERATE|>', pre_assistant.rstrip(' '))
+ result = result.replace('<|POST-ASSISTANT|>', post_assistant)
+
+ result = result.strip()
+
+ return result
diff --git a/modules/presets.py b/modules/presets.py
index b53195ee..560e0b77 100644
--- a/modules/presets.py
+++ b/modules/presets.py
@@ -16,9 +16,10 @@ default_preset_values = {
'dynatemp_exponent': 1,
'smoothing_factor': 0,
'smoothing_curve': 1,
- 'min_p': 0,
'top_p': 1,
'top_k': 0,
+ 'min_p': 0,
+ 'top_n_sigma': 0,
'typical_p': 1,
'xtc_threshold': 0.1,
'xtc_probability': 0,
@@ -26,7 +27,6 @@ default_preset_values = {
'eta_cutoff': 0,
'tfs': 1,
'top_a': 0,
- 'top_n_sigma': 0,
'adaptive_target': 0,
'adaptive_decay': 0.9,
'dry_multiplier': 0,
diff --git a/modules/prompts.py b/modules/prompts.py
index d107ce5a..85dc32e3 100644
--- a/modules/prompts.py
+++ b/modules/prompts.py
@@ -1,6 +1,7 @@
from pathlib import Path
from modules import shared, utils
+from modules.utils import sanitize_filename
from modules.text_generation import get_encoded_length
@@ -18,6 +19,7 @@ def load_prompt(fname):
return initial_content
+ fname = sanitize_filename(fname)
file_path = shared.user_data_dir / 'logs' / 'notebook' / f'{fname}.txt'
if file_path.exists():
with open(file_path, 'r', encoding='utf-8') as f:
diff --git a/modules/reasoning.py b/modules/reasoning.py
new file mode 100644
index 00000000..4a7cfa79
--- /dev/null
+++ b/modules/reasoning.py
@@ -0,0 +1,94 @@
+import html as html_module
+
+# Thinking block format definitions: (start_tag, end_tag, content_start_tag)
+# Use None for start_tag to match from beginning (end-only formats should be listed last)
+THINKING_FORMATS = [
+ ('', '', None),
+ ('<|channel|>analysis<|message|>', '<|end|>', '<|channel|>final<|message|>'),
+ ('<|channel|>commentary<|message|>', '<|end|>', '<|channel|>final<|message|>'),
+ ('', '', None),
+ ('<|channel>thought', '', None), # Gemma 4
+ ('<|think|>', '<|end|>', '<|content|>'), # Solar Open
+ # ('Thinking Process:', '', None), # Qwen3.5 verbose thinking outside tags -- removed: too prone to false positives in streaming
+ (None, '', None), # End-only variant (e.g., Qwen3-next)
+]
+
+
+def extract_reasoning(text, html_escaped=False):
+ """Extract reasoning/thinking blocks from the beginning of a string.
+
+ When html_escaped=True, tags are HTML-escaped before searching
+ (for use on already-escaped UI strings).
+
+ Returns (reasoning_content, final_content) where reasoning_content is
+ None if no thinking block is found.
+ """
+ if not text:
+ return None, text
+
+ esc = html_module.escape if html_escaped else lambda s: s
+
+ for start_tag, end_tag, content_tag in THINKING_FORMATS:
+ end_esc = esc(end_tag)
+ content_esc = esc(content_tag) if content_tag else None
+
+ if start_tag is None:
+ # End-only format: require end tag, start from beginning
+ end_pos = text.find(end_esc)
+ if end_pos == -1:
+ continue
+ thought_start = 0
+ else:
+ # Normal format: require start tag
+ start_esc = esc(start_tag)
+ start_pos = text.find(start_esc)
+ if start_pos == -1:
+ # During streaming, the start tag may be arriving partially.
+ # If the text is a prefix of a start tag, return empty content
+ # to prevent the partial tag from leaking.
+ stripped = text.strip()
+ if stripped and start_esc.startswith(stripped):
+ return '', ''
+ continue
+ thought_start = start_pos + len(start_esc)
+ end_pos = text.find(end_esc, thought_start)
+
+ if end_pos == -1:
+ # End tag missing - check if content tag can serve as fallback
+ if content_esc:
+ content_pos = text.find(content_esc, thought_start)
+ if content_pos != -1:
+ thought_end = content_pos
+ content_start = content_pos + len(content_esc)
+ else:
+ thought_end = len(text)
+ content_start = len(text)
+ else:
+ thought_end = len(text)
+ content_start = len(text)
+ else:
+ thought_end = end_pos
+ if content_esc:
+ content_pos = text.find(content_esc, end_pos)
+ if content_pos != -1:
+ content_start = content_pos + len(content_esc)
+ else:
+ # Content tag not present — fall back to content after
+ # end_tag (e.g. GPT-OSS tool calls skip the final channel).
+ content_start = end_pos + len(end_esc)
+ else:
+ content_start = end_pos + len(end_esc)
+
+ return text[thought_start:thought_end], text[content_start:].lstrip()
+
+ # Handle standalone GPT-OSS final channel marker without a preceding
+ # analysis/commentary block (the model skipped thinking entirely).
+ for marker in ['<|start|>assistant<|channel|>final<|message|>', '<|channel|>final<|message|>']:
+ marker_esc = esc(marker)
+ pos = text.find(marker_esc)
+ if pos != -1:
+ before = text[:pos].strip()
+ after = text[pos + len(marker_esc):]
+ return (before if before else None), after
+
+ return None, text
diff --git a/modules/shared.py b/modules/shared.py
index bc7ea8ba..13843f0c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -47,7 +47,7 @@ parser = argparse.ArgumentParser(description="Text Generation Web UI", conflict_
# Basic settings
group = parser.add_argument_group('Basic settings')
group.add_argument('--user-data-dir', type=str, default=str(user_data_dir), help='Path to the user data directory. Default: auto-detected.')
-group.add_argument('--multi-user', action='store_true', help='Multi-user mode. Chat histories are not saved or automatically loaded. Warning: this is likely not safe for sharing publicly.')
+group.add_argument('--multi-user', action='store_true', help='Multi-user mode. Chat histories are not saved or automatically loaded. Best suited for small trusted teams.')
group.add_argument('--model', type=str, help='Name of the model to load by default.')
group.add_argument('--lora', type=str, nargs='+', help='The list of LoRAs to load. If you want to load more than one LoRA, write the names separated by spaces.')
group.add_argument('--model-dir', type=str, default=str(user_data_dir / 'models'), help='Path to directory with all the models.')
@@ -76,7 +76,7 @@ group.add_argument('--loader', type=str, help='Choose the model loader manually,
# Cache
group = parser.add_argument_group('Context and cache')
-group.add_argument('--ctx-size', '--n_ctx', '--max_seq_len', type=int, default=8192, metavar='N', help='Context size in tokens. llama.cpp: 0 = auto if gpu-layers is also -1.')
+group.add_argument('--ctx-size', '--n_ctx', '--max_seq_len', type=int, default=0, metavar='N', help='Context size in tokens. 0 = auto for llama.cpp (requires gpu-layers=-1), 8192 for other loaders.')
group.add_argument('--cache-type', '--cache_type', type=str, default='fp16', metavar='N', help='KV cache type; valid options: llama.cpp - fp16, q8_0, q4_0; ExLlamaV3 - fp16, q2 to q8 (can specify k_bits and v_bits separately, e.g. q4_q8).')
# Speculative decoding
@@ -101,15 +101,16 @@ group.add_argument('--tensor-split', type=str, default=None, help='Split the mod
group.add_argument('--row-split', action='store_true', help='Split the model by rows across GPUs. This may improve multi-gpu performance.')
group.add_argument('--no-mmap', action='store_true', help='Prevent mmap from being used.')
group.add_argument('--mlock', action='store_true', help='Force the system to keep the model in RAM.')
-group.add_argument('--no-kv-offload', action='store_true', help='Do not offload the K, Q, V to the GPU. This saves VRAM but reduces the performance.')
+group.add_argument('--no-kv-offload', action='store_true', help='Do not offload the K, Q, V to the GPU. This saves VRAM but reduces performance.')
group.add_argument('--batch-size', type=int, default=1024, help='Maximum number of prompt tokens to batch together when calling llama-server. This is the application level batch size.')
group.add_argument('--ubatch-size', type=int, default=1024, help='Maximum number of prompt tokens to batch together when calling llama-server. This is the max physical batch size for computation (device level).')
group.add_argument('--threads', type=int, default=0, help='Number of threads to use.')
group.add_argument('--threads-batch', type=int, default=0, help='Number of threads to use for batches/prompt processing.')
group.add_argument('--numa', action='store_true', help='Activate NUMA task allocation for llama.cpp.')
group.add_argument('--parallel', type=int, default=1, help='Number of parallel request slots. The context size is divided equally among slots. For example, to have 4 slots with 8192 context each, set ctx_size to 32768.')
-group.add_argument('--fit-target', type=str, default='1024', help='Target VRAM margin per device for auto GPU layers, comma-separated list of values in MiB. A single value is broadcast across all devices. Default: 1024.')
-group.add_argument('--extra-flags', type=str, default=None, help='Extra flags to pass to llama-server. Format: "flag1=value1,flag2,flag3=value3". Example: "override-tensor=exps=CPU"')
+group.add_argument('--fit-target', type=str, default='512', help='Target VRAM margin per device for auto GPU layers, comma-separated list of values in MiB. A single value is broadcast across all devices.')
+group.add_argument('--extra-flags', type=str, default=None, help='Extra flags to pass to llama-server. Example: "--jinja --rpc 192.168.1.100:50052"')
+group.add_argument('--ik', action='store_true', help='Use ik_llama.cpp instead of upstream llama.cpp. Requires the ik_llama_cpp_binaries package to be installed.')
# Transformers/Accelerate
group = parser.add_argument_group('Transformers/Accelerate')
@@ -139,12 +140,6 @@ group.add_argument('--enable-tp', '--enable_tp', action='store_true', help='Enab
group.add_argument('--tp-backend', type=str, default='native', help='The backend for tensor parallelism. Valid options: native, nccl. Default: native.')
group.add_argument('--cfg-cache', action='store_true', help='Create an additional cache for CFG negative prompts. Necessary to use CFG with that loader.')
-# RoPE
-group = parser.add_argument_group('RoPE')
-group.add_argument('--alpha_value', type=float, default=1, help='Positional embeddings alpha factor for NTK RoPE scaling. Use either this or compress_pos_emb, not both.')
-group.add_argument('--rope_freq_base', type=int, default=0, help='If greater than 0, will be used instead of alpha_value. Those two are related by rope_freq_base = 10000 * alpha_value ^ (64 / 63).')
-group.add_argument('--compress_pos_emb', type=int, default=1, help="Positional embeddings compression factor. Should be set to (context length) / (model\'s original context length). Equal to 1/rope_freq_scale.")
-
# Gradio
group = parser.add_argument_group('Gradio')
group.add_argument('--listen', action='store_true', help='Make the web UI reachable from your local network.')
@@ -162,8 +157,8 @@ group.add_argument('--portable', action='store_true', help='Hide features not av
# API
group = parser.add_argument_group('API')
-group.add_argument('--api', action='store_true', help='Enable the API extension.')
-group.add_argument('--public-api', action='store_true', help='Create a public URL for the API using Cloudfare.')
+group.add_argument('--api', action='store_true', help='Enable the API server.')
+group.add_argument('--public-api', action='store_true', help='Create a public URL for the API using Cloudflare.')
group.add_argument('--public-api-id', type=str, help='Tunnel ID for named Cloudflare Tunnel. Use together with public-api option.', default=None)
group.add_argument('--api-port', type=int, default=5000, help='The listening port for the API.')
group.add_argument('--api-key', type=str, default='', help='API authentication key.')
@@ -181,9 +176,10 @@ group.add_argument('--dynatemp-high', type=float, default=_d['dynatemp_high'], m
group.add_argument('--dynatemp-exponent', type=float, default=_d['dynatemp_exponent'], metavar='N', help='Dynamic temperature exponent')
group.add_argument('--smoothing-factor', type=float, default=_d['smoothing_factor'], metavar='N', help='Smoothing factor')
group.add_argument('--smoothing-curve', type=float, default=_d['smoothing_curve'], metavar='N', help='Smoothing curve')
-group.add_argument('--min-p', type=float, default=_d['min_p'], metavar='N', help='Min P')
-group.add_argument('--top-p', type=float, default=_d['top_p'], metavar='N', help='Top P')
+group.add_argument('--top-p', type=float, default=0.95, metavar='N', help='Top P')
group.add_argument('--top-k', type=int, default=_d['top_k'], metavar='N', help='Top K')
+group.add_argument('--min-p', type=float, default=_d['min_p'], metavar='N', help='Min P')
+group.add_argument('--top-n-sigma', type=float, default=_d['top_n_sigma'], metavar='N', help='Top N Sigma')
group.add_argument('--typical-p', type=float, default=_d['typical_p'], metavar='N', help='Typical P')
group.add_argument('--xtc-threshold', type=float, default=_d['xtc_threshold'], metavar='N', help='XTC threshold')
group.add_argument('--xtc-probability', type=float, default=_d['xtc_probability'], metavar='N', help='XTC probability')
@@ -191,7 +187,6 @@ group.add_argument('--epsilon-cutoff', type=float, default=_d['epsilon_cutoff'],
group.add_argument('--eta-cutoff', type=float, default=_d['eta_cutoff'], metavar='N', help='Eta cutoff')
group.add_argument('--tfs', type=float, default=_d['tfs'], metavar='N', help='TFS')
group.add_argument('--top-a', type=float, default=_d['top_a'], metavar='N', help='Top A')
-group.add_argument('--top-n-sigma', type=float, default=_d['top_n_sigma'], metavar='N', help='Top N Sigma')
group.add_argument('--adaptive-target', type=float, default=_d['adaptive_target'], metavar='N', help='Adaptive target')
group.add_argument('--adaptive-decay', type=float, default=_d['adaptive_decay'], metavar='N', help='Adaptive decay')
group.add_argument('--dry-multiplier', type=float, default=_d['dry_multiplier'], metavar='N', help='DRY multiplier')
@@ -263,8 +258,9 @@ settings = {
'chat-instruct_command': 'Continue the chat dialogue below. Write a single reply for the character "<|character|>". Reply directly, without starting the reply with the character name.\n\n<|prompt|>',
'enable_web_search': False,
'web_search_pages': 3,
+ 'selected_tools': [],
'prompt-notebook': '',
- 'preset': 'Qwen3 - Thinking' if (user_data_dir / 'presets/Qwen3 - Thinking.yaml').exists() else None,
+ 'preset': 'Top-P' if (user_data_dir / 'presets/Top-P.yaml').exists() else None,
'max_new_tokens': 512,
'max_new_tokens_min': 1,
'max_new_tokens_max': 4096,
@@ -289,7 +285,7 @@ settings = {
'include_past_attachments': True,
# Generation parameters - Curve shape
- 'temperature': 0.6,
+ 'temperature': neutral_samplers['temperature'],
'dynatemp_low': neutral_samplers['dynatemp_low'],
'dynatemp_high': neutral_samplers['dynatemp_high'],
'dynatemp_exponent': neutral_samplers['dynatemp_exponent'],
@@ -297,9 +293,10 @@ settings = {
'smoothing_curve': neutral_samplers['smoothing_curve'],
# Generation parameters - Curve cutoff
- 'min_p': neutral_samplers['min_p'],
'top_p': 0.95,
- 'top_k': 20,
+ 'top_k': neutral_samplers['top_k'],
+ 'min_p': neutral_samplers['min_p'],
+ 'top_n_sigma': neutral_samplers['top_n_sigma'],
'typical_p': neutral_samplers['typical_p'],
'xtc_threshold': neutral_samplers['xtc_threshold'],
'xtc_probability': neutral_samplers['xtc_probability'],
@@ -307,7 +304,6 @@ settings = {
'eta_cutoff': neutral_samplers['eta_cutoff'],
'tfs': neutral_samplers['tfs'],
'top_a': neutral_samplers['top_a'],
- 'top_n_sigma': neutral_samplers['top_n_sigma'],
'adaptive_target': neutral_samplers['adaptive_target'],
'adaptive_decay': neutral_samplers['adaptive_decay'],
@@ -347,7 +343,7 @@ settings = {
'greeting': 'How can I help you today?',
'custom_system_message': '',
'instruction_template_str': "{%- set ns = namespace(found=false) -%}\n{%- for message in messages -%}\n {%- if message['role'] == 'system' -%}\n {%- set ns.found = true -%}\n {%- endif -%}\n{%- endfor -%}\n{%- if not ns.found -%}\n {{- '' + 'Below is an instruction that describes a task. Write a response that appropriately completes the request.' + '\\n\\n' -}}\n{%- endif %}\n{%- for message in messages %}\n {%- if message['role'] == 'system' -%}\n {{- '' + message['content'] + '\\n\\n' -}}\n {%- else -%}\n {%- if message['role'] == 'user' -%}\n {{-'### Instruction:\\n' + message['content'] + '\\n\\n'-}}\n {%- else -%}\n {{-'### Response:\\n' + message['content'] + '\\n\\n' -}}\n {%- endif -%}\n {%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt -%}\n {{-'### Response:\\n'-}}\n{%- endif -%}",
- 'chat_template_str': "{%- for message in messages %}\n {%- if message['role'] == 'system' -%}\n {%- if message['content'] -%}\n {{- message['content'] + '\\n\\n' -}}\n {%- endif -%}\n {%- if user_bio -%}\n {{- user_bio + '\\n\\n' -}}\n {%- endif -%}\n {%- else -%}\n {%- if message['role'] == 'user' -%}\n {{- name1 + ': ' + message['content'] + '\\n'-}}\n {%- else -%}\n {{- name2 + ': ' + message['content'] + '\\n' -}}\n {%- endif -%}\n {%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt %}\n {{- name2 + ':' -}}\n{%- endif %}",
+ 'chat_template_str': "{%- for message in messages %}\n {%- if message['role'] == 'system' -%}\n {%- if message['content'] -%}\n {{- message['content'] + '\\n\\n' -}}\n {%- endif -%}\n {%- if user_bio -%}\n {{- user_bio + '\\n\\n' -}}\n {%- endif -%}\n {%- elif message['role'] == 'tool' -%}\n {{- '[Tool result: ' + message['content'] + ']\\n' -}}\n {%- elif message['role'] == 'user' -%}\n {{- name1 + ': ' + message['content'] + '\\n'-}}\n {%- elif message['tool_calls'] is defined and message['tool_calls'] -%}\n {%- for tc in message['tool_calls'] -%}\n {{- '[Calling: ' + tc['function']['name'] + '(' + tc['function']['arguments'] + ')]\\n' -}}\n {%- endfor -%}\n {%- else -%}\n {{- name2 + ': ' + message['content'] + '\\n' -}}\n {%- endif -%}\n{%- endfor -%}\n{%- if add_generation_prompt %}\n {{- name2 + ':' -}}\n{%- endif %}",
# Extensions
'default_extensions': [],
@@ -395,9 +391,16 @@ def do_cmd_flags_warnings():
if args.share:
logger.warning("The gradio \"share link\" feature uses a proprietary executable to create a reverse tunnel. Use it with care.")
if any((args.listen, args.share)) and not any((args.gradio_auth, args.gradio_auth_path)):
- logger.warning("\nYou are potentially exposing the web UI to the entire internet without any access password.\nYou can create one with the \"--gradio-auth\" flag like this:\n\n--gradio-auth username:password\n\nMake sure to replace username:password with your own.")
- if args.multi_user:
- logger.warning('\nThe multi-user mode is highly experimental and should not be shared publicly.')
+ logger.warning("You are potentially exposing the web UI to the entire internet without any access password.\nYou can create one with the \"--gradio-auth\" flag like this:\n\n--gradio-auth username:password\n\nMake sure to replace username:password with your own.")
+ if args.multi_user:
+ logger.warning(
+ 'Multi-user mode is enabled. Known limitations:'
+ '\n- The Stop button stops generation for all users, not just you.'
+ '\n- Chat history is not saved and will be lost on page refresh.'
+ '\n- Only one user can generate at a time unless using a parallel-capable backend (e.g. llama.cpp with --parallel N for N > 1, or ExLlamaV3).'
+ '\n\nThis mode works best for small trusted teams.'
+ '\n\nDo not expose publicly. Grayed-out actions can easily be bypassed client-side.\n'
+ )
def apply_image_model_cli_overrides():
@@ -433,16 +436,6 @@ def fix_loader_name(name):
return 'TensorRT-LLM'
-def add_extension(name, last=False):
- if args.extensions is None:
- args.extensions = [name]
- elif last:
- args.extensions = [x for x in args.extensions if x != name]
- args.extensions.append(name)
- elif name not in args.extensions:
- args.extensions.append(name)
-
-
def is_chat():
return True
@@ -451,36 +444,18 @@ def load_user_config():
'''
Loads custom model-specific settings
'''
+ user_config = {}
if Path(f'{args.model_dir}/config-user.yaml').exists():
file_content = open(f'{args.model_dir}/config-user.yaml', 'r').read().strip()
-
if file_content:
user_config = yaml.safe_load(file_content)
- else:
- user_config = {}
- else:
- user_config = {}
return user_config
args.loader = fix_loader_name(args.loader)
-# Activate the API extension
-if args.api or args.public_api:
- add_extension('openai', last=True)
-
-# Load model-specific settings
-p = Path(f'{args.model_dir}/config.yaml')
-if p.exists():
- model_config = yaml.safe_load(open(p, 'r').read())
-else:
- model_config = {}
-del p
-
-
# Load custom model-specific settings
user_config = load_user_config()
-model_config = OrderedDict(model_config)
user_config = OrderedDict(user_config)
diff --git a/modules/text_generation.py b/modules/text_generation.py
index c78afe3e..3a9ddab5 100644
--- a/modules/text_generation.py
+++ b/modules/text_generation.py
@@ -4,7 +4,6 @@ import html
import pprint
import random
import time
-import traceback
import numpy as np
@@ -18,9 +17,7 @@ from modules.utils import check_model_loaded
def generate_reply(*args, **kwargs):
- if shared.args.idle_timeout > 0 and shared.model is None and shared.model_name not in [None, 'None']:
- from modules.models import load_model
- shared.model, shared.tokenizer = load_model(shared.model_name)
+ models.load_model_if_idle_unloaded()
state = args[1] if len(args) > 1 else kwargs.get('state', {})
use_parallel = (
@@ -32,10 +29,16 @@ def generate_reply(*args, **kwargs):
if not use_parallel:
shared.generation_lock.acquire()
+ with models._generation_count_lock:
+ models.active_generation_count += 1
+
try:
for result in _generate_reply(*args, **kwargs):
yield result
finally:
+ with models._generation_count_lock:
+ models.active_generation_count -= 1
+
models.last_generation_time = time.time()
if not use_parallel:
shared.generation_lock.release()
@@ -78,10 +81,13 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
reply = ''
is_stream = state['stream']
if len(all_stop_strings) > 0 and not state['stream']:
+ original_logits_processor = state.get('logits_processor')
stop_event_ref = state.pop('stop_event', None)
state = copy.deepcopy(state)
if stop_event_ref is not None:
state['stop_event'] = stop_event_ref
+ if original_logits_processor is not None:
+ state['logits_processor'] = original_logits_processor
state['stream'] = True
# Generate
@@ -124,7 +130,9 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False, escap
def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_length=None):
if shared.tokenizer is None:
- raise ValueError('No tokenizer is loaded')
+ models.load_model_if_idle_unloaded()
+ if shared.tokenizer is None:
+ raise ValueError('No tokenizer is loaded')
# llama.cpp case
if shared.model.__class__.__name__ == 'LlamaServer':
@@ -174,7 +182,9 @@ def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_lengt
def decode(output_ids, skip_special_tokens=True):
if shared.tokenizer is None:
- raise ValueError('No tokenizer is loaded')
+ models.load_model_if_idle_unloaded()
+ if shared.tokenizer is None:
+ raise ValueError('No tokenizer is loaded')
return shared.tokenizer.decode(output_ids, skip_special_tokens=skip_special_tokens)
@@ -375,7 +385,7 @@ def generate_reply_HF(question, original_question, state, stopping_strings=None,
generate_params['sampler_priority'] = [x.strip() for x in state['sampler_priority'].replace('\n', ',').split(',') if x.strip()]
if state['custom_token_bans']:
- to_ban = [int(x) for x in state['custom_token_bans'].split(',')]
+ to_ban = [int(x.strip()) for x in state['custom_token_bans'].split(',') if x.strip()]
if len(to_ban) > 0:
if generate_params.get('suppress_tokens', None):
generate_params['suppress_tokens'] += to_ban
@@ -474,7 +484,7 @@ def generate_reply_HF(question, original_question, state, stopping_strings=None,
yield cumulative_reply
except Exception:
- traceback.print_exc()
+ logger.exception("Failed to generate reply (HF)")
finally:
t1 = time.time()
original_tokens = len(original_input_ids[0])
@@ -507,7 +517,7 @@ def generate_reply_custom(question, original_question, state, stopping_strings=N
yield reply
except Exception:
- traceback.print_exc()
+ logger.exception("Failed to generate reply (custom)")
finally:
t1 = time.time()
diff --git a/modules/tool_parsing.py b/modules/tool_parsing.py
new file mode 100644
index 00000000..919e523a
--- /dev/null
+++ b/modules/tool_parsing.py
@@ -0,0 +1,711 @@
+import json
+import random
+import re
+
+from modules.reasoning import extract_reasoning
+
+
+def _make_tool_call(name, arguments):
+ return {"type": "function", "function": {"name": name, "arguments": arguments}}
+
+
+def get_tool_call_id() -> str:
+ letter_bytes = "abcdefghijklmnopqrstuvwxyz0123456789"
+ b = [random.choice(letter_bytes) for _ in range(8)]
+ return "call_" + "".join(b).lower()
+
+
+# All known opening markers for tool calls across model formats.
+TOOL_CALL_OPENING_MARKERS = [
+ '',
+ '',
+ '',
+ '<|tool_call_begin|>',
+ '<|tool_calls_section_begin|>',
+ '<｜tool▁call▁begin｜>',
+ '<｜tool▁calls▁begin｜>',
+ '[TOOL_CALLS]',
+ 'to=functions.',
+ '<|channel|>commentary',
+ '<|tool_call>call:',
+]
+
+
+def streaming_tool_buffer_check(text, markers=None, tool_names=None, check_bare_names=False):
+ '''
+ Check whether streaming output should be withheld because it may
+ contain tool-call markup.
+
+ Args:
+ text: Full accumulated internal text.
+ markers: Template-specific markers for partial-prefix matching.
+ If None, falls back to TOOL_CALL_OPENING_MARKERS.
+ tool_names: List of tool function names.
+ check_bare_names: Whether to do partial-prefix matching on tool
+ names (for models with unknown template format).
+ '''
+ # Strip thinking blocks so tool-call syntax inside doesn't
+ # trigger false positives.
+ _, text = extract_reasoning(text)
+
+ # Full marker found in text → buffer permanently.
+ # Always checks ALL known markers regardless of template (cheap safety net).
+ for marker in TOOL_CALL_OPENING_MARKERS:
+ if marker in text:
+ return True
+
+ # Bare function-name full match: "get_weather{...}" or "get_weather {...}"
+ if tool_names:
+ for name in tool_names:
+ if name + '{' in text or name + ' {' in text:
+ return True
+
+ # Partial-prefix matching: only for template-specific markers.
+ for marker in (markers if markers is not None else TOOL_CALL_OPENING_MARKERS):
+ for prefix_len in range(min(len(marker) - 1, len(text)), 0, -1):
+ if text.endswith(marker[:prefix_len]):
+ return True
+
+ # Bare-name partial matching: only when template format is unknown.
+ if check_bare_names and tool_names:
+ for name in tool_names:
+ if text.endswith(name):
+ return True
+ for prefix_len in range(min(len(name) - 1, len(text)), 0, -1):
+ if text.endswith(name[:prefix_len]):
+ return True
+
+ return False
+
+
+def check_and_sanitize_tool_call_candidate(candidate_dict: dict, tool_names: list[str]):
+ # check if property 'function' exists and is a dictionary, otherwise adapt dict
+ if 'function' not in candidate_dict and 'name' in candidate_dict and isinstance(candidate_dict['name'], str):
+ candidate_dict = {"type": "function", "function": candidate_dict}
+ if 'function' in candidate_dict and isinstance(candidate_dict['function'], str):
+ candidate_dict['name'] = candidate_dict['function']
+ del candidate_dict['function']
+ candidate_dict = {"type": "function", "function": candidate_dict}
+ if 'function' in candidate_dict and isinstance(candidate_dict['function'], dict):
+ # check if 'name' exists within 'function' and is part of known tools
+ if 'name' in candidate_dict['function'] and candidate_dict['function']['name'] in tool_names:
+ candidate_dict["type"] = "function" # ensure required property 'type' exists and has the right value
+ # map property 'parameters' used by some older models to 'arguments'
+ if "arguments" not in candidate_dict["function"] and "parameters" in candidate_dict["function"]:
+ candidate_dict["function"]["arguments"] = candidate_dict["function"]["parameters"]
+ del candidate_dict["function"]["parameters"]
+ return candidate_dict
+ return None
+
+
+def _extract_balanced_json(text: str, start: int) -> str | None:
+ """Extract a balanced JSON object from text starting at the given position.
+
+ Walks through the string tracking brace depth and string boundaries
+ to correctly handle arbitrary nesting levels.
+ """
+ if start >= len(text) or text[start] != '{':
+ return None
+ depth = 0
+ in_string = False
+ escape_next = False
+ for i in range(start, len(text)):
+ c = text[i]
+ if escape_next:
+ escape_next = False
+ continue
+ if c == '\\' and in_string:
+ escape_next = True
+ continue
+ if c == '"':
+ in_string = not in_string
+ continue
+ if in_string:
+ continue
+ if c == '{':
+ depth += 1
+ elif c == '}':
+ depth -= 1
+ if depth == 0:
+ return text[start:i + 1]
+ return None
+
+
+def _parse_channel_tool_calls(answer: str, tool_names: list[str]):
+ """Parse channel-based tool calls used by GPT-OSS and similar models.
+
+ Format:
+ <|start|>assistant to=functions.func_name<|channel|>commentary json<|message|>{"arg": "value"}
+ or:
+ <|channel|>commentary to=functions.func_name <|constrain|>json<|message|>{"arg": "value"}
+ """
+ matches = []
+ start_pos = None
+ # Pattern 1: to=functions.NAME before <|channel|> (GPT-OSS primary format)
+ # Pattern 2: to=functions.NAME after <|channel|> (alternative format)
+ patterns = [
+ r'to=functions\.([^<\s]+)\s*<\|channel\|>[^<]*<\|message\|>',
+ r'<\|channel\|>\w+ to=functions\.([^<\s]+).*?<\|message\|>',
+ ]
+ for pattern in patterns:
+ for m in re.finditer(pattern, answer):
+ func_name = m.group(1).strip()
+ if func_name not in tool_names:
+ continue
+ json_str = _extract_balanced_json(answer, m.end())
+ if json_str is None:
+ continue
+ try:
+ arguments = json.loads(json_str)
+ if start_pos is None:
+ prefix = answer.rfind('<|start|>assistant', 0, m.start())
+ start_pos = prefix if prefix != -1 else m.start()
+ matches.append(_make_tool_call(func_name, arguments))
+ except json.JSONDecodeError:
+ pass
+ if matches:
+ break
+ return matches, start_pos
+
+
+def _parse_mistral_token_tool_calls(answer: str, tool_names: list[str]):
+ """Parse Mistral/Devstral-style tool calls with [TOOL_CALLS] and [ARGS] special tokens.
+
+ Format:
+ [TOOL_CALLS]func_name[ARGS]{"arg": "value"}
+ """
+ matches = []
+ start_pos = None
+ for m in re.finditer(
+ r'\[TOOL_CALLS\]\s*(\S+?)\s*\[ARGS\]\s*',
+ answer
+ ):
+ func_name = m.group(1).strip()
+ if func_name not in tool_names:
+ continue
+ json_str = _extract_balanced_json(answer, m.end())
+ if json_str is None:
+ continue
+ try:
+ arguments = json.loads(json_str)
+ if start_pos is None:
+ start_pos = m.start()
+ matches.append(_make_tool_call(func_name, arguments))
+ except json.JSONDecodeError:
+ pass
+ return matches, start_pos
+
+
+def _parse_bare_name_tool_calls(answer: str, tool_names: list[str]):
+ """Parse bare function-name style tool calls used by Mistral and similar models.
+
+ Format:
+ functionName{"arg": "value"}
+ Multiple calls are concatenated directly or separated by whitespace.
+ """
+ matches = []
+ start_pos = None
+ # Match tool name followed by opening brace, then extract balanced JSON
+ escaped_names = [re.escape(name) for name in tool_names]
+ pattern = r'(?:' + '|'.join(escaped_names) + r')\s*\{'
+ for match in re.finditer(pattern, answer):
+ text = match.group(0)
+ name = None
+ for n in tool_names:
+ if text.startswith(n):
+ name = n
+ break
+ if not name:
+ continue
+ brace_start = match.end() - 1
+ json_str = _extract_balanced_json(answer, brace_start)
+ if json_str is None:
+ continue
+ try:
+ arguments = json.loads(json_str)
+ if start_pos is None:
+ start_pos = match.start()
+ matches.append(_make_tool_call(name, arguments))
+ except json.JSONDecodeError:
+ pass
+ return matches, start_pos
+
+
+def _parse_xml_param_tool_calls(answer: str, tool_names: list[str]):
+ """Parse XML-parameter style tool calls used by Qwen3.5 and similar models.
+
+ Format:
+
+
+ value
+