File size: 4,245 Bytes
45b2cda 9f88ad1 6b23da9 9f88ad1 6b23da9 1f6095e 6b23da9 1f6095e 6b23da9 45b2cda 1f6095e fd3d17e 1f6095e 45b2cda 1f6095e 6b23da9 1f6095e 6b23da9 9f88ad1 6b23da9 fd3d17e 6b23da9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | """Tinker inference client. Supports both base models and fine-tuned checkpoints."""
import json
import re
import streamlit as st
def _should_print_model_input(cfg: dict) -> bool:
"""True if user enabled PRINT_MODEL_INPUT env or print_model_input in study_config.yaml."""
return bool(cfg.get("print_model_input", False))
def _print_model_input(messages: list, prompt: object) -> None:
"""Dump logical messages + exact string passed to sampling_client.sample(prompt=...)."""
slim = [{"role": m.get("role"), "content": m.get("content")} for m in messages]
sep = "=" * 72
print(f"\n{sep}\nMODEL INPUT — logical messages (role/content)\n{sep}", flush=True)
print(json.dumps(slim, indent=2, ensure_ascii=False), flush=True)
prompt_str = prompt if isinstance(prompt, str) else str(prompt)
print(f"\n{sep}\nMODEL INPUT — rendered prompt (build_generation_prompt)\n{sep}", flush=True)
print(prompt_str, flush=True)
print(f"{sep}\nEND MODEL INPUT\n", flush=True)
@st.cache_resource
def _get_tinker_clients(model_name: str, sampler_path: str = ""):
"""
Initialise and cache the Tinker sampling client, renderer, and tokenizer.
If sampler_path is provided, loads from that checkpoint (fine-tuned model).
Otherwise, loads the base model_name.
Cache key includes both so different variants get different clients.
"""
import tinker
from tinker import types as tinker_types
from tinker_cookbook import renderers
from tinker_cookbook.model_info import get_recommended_renderer_name
from tinker_cookbook.tokenizer_utils import get_tokenizer
service_client = tinker.ServiceClient()
if sampler_path:
print(f"[MODEL] Loading fine-tuned checkpoint: {sampler_path}")
sampling_client = service_client.create_sampling_client(model_path=sampler_path)
else:
print(f"[MODEL] Loading base model: {model_name}")
sampling_client = service_client.create_sampling_client(base_model=model_name)
tokenizer = get_tokenizer(model_name)
renderer_name = get_recommended_renderer_name(model_name)
renderer = renderers.get_renderer(renderer_name, tokenizer)
return sampling_client, renderer, tinker_types
def call_model(messages: list, cfg: dict) -> str:
"""Send a message list to Tinker and return cleaned response text."""
model_name = cfg["model_name"]
sampler_path = cfg.get("sampler_path", "")
temperature = float(cfg.get("sampling_temperature", 1.0))
print(
f"[MODEL] model_name={model_name} sampler_path={sampler_path or '(base)'} "
f"temperature={temperature}"
)
print(f"[MODEL] num_messages={len(messages)}")
print(f"[MODEL] roles={[m['role'] for m in messages]}")
if messages:
print(f"[MODEL] system_prompt[:150]={messages[0]['content'][:150]}")
try:
from tinker_cookbook import renderers as tinker_renderers
sampling_client, renderer, tinker_types = _get_tinker_clients(model_name, sampler_path)
prompt = renderer.build_generation_prompt(messages)
if _should_print_model_input(cfg):
_print_model_input(messages, prompt)
params = tinker_types.SamplingParams(
max_tokens=1000,
temperature=temperature,
stop=renderer.get_stop_sequences(),
)
result = sampling_client.sample(
prompt=prompt,
sampling_params=params,
num_samples=1,
).result()
parsed_message, _ = renderer.parse_response(result.sequences[0].tokens)
content = tinker_renderers.format_content_as_string(parsed_message["content"])
content = re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL).strip()
content = re.sub(r"<\|[^|]*\|>", "", content).strip()
match = re.search(r"(.{40,}?)\1{4,}", content, flags=re.DOTALL)
if match:
content = content[: match.start() + len(match.group(1))].strip()
if not content or len(content.split()) < 3:
raise ValueError("Model output cleanup yielded no usable content.")
return content
except Exception as e:
print(f"[MODEL] Tinker error: {e}")
return f"[Model error: {e}]" |