mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
580 lines
18 KiB
Python
580 lines
18 KiB
Python
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import gradio as gr
|
|
import torch
|
|
|
|
# Add the src directory to the path
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
|
|
|
from ctx_to_lora.data.processing import tokenize_ctx_text
|
|
from ctx_to_lora.model_loading import get_tokenizer
|
|
from ctx_to_lora.modeling import hypernet
|
|
|
|
sys.modules["ctx_to_lora.modeling_utils"] = hypernet
|
|
|
|
# Global state
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
modulated_model = None
|
|
chat_history = []
|
|
ctx_tokenizer = None
|
|
base_tokenizer = None
|
|
|
|
try:
|
|
DEFAULT_CONTEXT = Path("data/sakana_wiki.txt").read_text(encoding="utf-8").strip()
|
|
except FileNotFoundError:
|
|
DEFAULT_CONTEXT = ""
|
|
|
|
WARNING_MESSAGE = (
|
|
"⚠️ **Caution**: This is an academic proof-of-concept application.\n"
|
|
"The model may generate inaccurate information or hallucinate facts."
|
|
)
|
|
|
|
FOOTER = """
|
|
⚠️ This model is an experimental prototype and is only available for educational and research and development purposes. It is not suitable for commercial use or in environments where failure can have significant effects (mission-critical environments).
|
|
The use of this model is at the user's own risk and its performance and results is not guaranteed in any way.
|
|
Sakana AI is not responsible for any direct or indirect loss resulting from using this model, regardless of the outcome.
|
|
"""
|
|
|
|
|
|
def load_custom_chat_template(tokenizer, model_name):
|
|
if "gemma" in model_name.lower():
|
|
template_path = "chat_templates/google/gemma-2-2b-it.jinja"
|
|
if os.path.exists(template_path):
|
|
with open(template_path) as f:
|
|
template_content = f.read()
|
|
tokenizer.chat_template = template_content
|
|
print(f"Loaded custom chat template from {template_path}")
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_available_checkpoints():
|
|
checkpoints = []
|
|
trained_dir = Path("trained_t2l")
|
|
|
|
if trained_dir.exists():
|
|
for model_dir in trained_dir.iterdir():
|
|
if model_dir.is_dir():
|
|
for ckpt_file in model_dir.glob("**/*.bin"):
|
|
if "pytorch_model.bin" in ckpt_file.name:
|
|
checkpoints.append(str(ckpt_file))
|
|
|
|
train_outputs = Path("train_outputs/runs")
|
|
if train_outputs.exists():
|
|
for run_dir in train_outputs.iterdir():
|
|
if run_dir.is_dir():
|
|
for ckpt_dir in run_dir.glob("checkpoint-*"):
|
|
ckpt_file = ckpt_dir / "pytorch_model.bin"
|
|
if ckpt_file.exists():
|
|
checkpoints.append(str(ckpt_file))
|
|
|
|
return sorted(checkpoints) if checkpoints else ["No checkpoints found"]
|
|
|
|
|
|
def load_checkpoint(checkpoint_path: str) -> tuple[str, gr.update]:
|
|
global modulated_model, ctx_tokenizer, base_tokenizer, chat_history
|
|
|
|
if not checkpoint_path or checkpoint_path == "No checkpoints found":
|
|
return "⚠️ Please select a valid checkpoint", gr.update()
|
|
|
|
try:
|
|
print(f"Loading checkpoint: {checkpoint_path}")
|
|
from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
|
|
|
state_dict = torch.load(checkpoint_path, weights_only=False)
|
|
modulated_model = ModulatedPretrainedModel.from_state_dict(
|
|
state_dict,
|
|
train=False,
|
|
use_flash_attn=True,
|
|
use_sequence_packing=False,
|
|
)
|
|
modulated_model = modulated_model.to(device).to(torch.bfloat16)
|
|
modulated_model.eval()
|
|
|
|
ctx_encoder_model_name_or_path = (
|
|
modulated_model.ctx_encoder_args.ctx_encoder_model_name_or_path
|
|
or modulated_model.base_model.config.name_or_path
|
|
)
|
|
ctx_tokenizer = get_tokenizer(ctx_encoder_model_name_or_path, train=False)
|
|
base_tokenizer = get_tokenizer(
|
|
modulated_model.base_model.config.name_or_path, train=False
|
|
)
|
|
|
|
load_custom_chat_template(
|
|
base_tokenizer, modulated_model.base_model.config.name_or_path
|
|
)
|
|
|
|
chat_history = [{"role": "system", "content": ""}]
|
|
|
|
model_name = modulated_model.base_model.config.name_or_path
|
|
success_msg = (
|
|
f"✅ Successfully loaded checkpoint!\n\nBase Model: {model_name}\n\n"
|
|
"You can now add context and start chatting."
|
|
)
|
|
return success_msg, gr.update(interactive=True)
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = (
|
|
f"❌ Error loading checkpoint:\n{str(e)}\n\n{traceback.format_exc()}"
|
|
)
|
|
print(error_msg)
|
|
return error_msg, gr.update(interactive=False)
|
|
|
|
|
|
def process_multiple_contexts(contexts: list[str]) -> dict:
|
|
contexts = [ctx for ctx in contexts if ctx.strip()]
|
|
if not contexts:
|
|
contexts = [""]
|
|
print(f"Processing {len(contexts)} non-empty contexts")
|
|
tokenized_contexts = tokenize_ctx_text({"context": contexts}, ctx_tokenizer)
|
|
ctx_ids = tokenized_contexts["ctx_ids"]
|
|
ctx_ids = [
|
|
torch.tensor(ctx_id, dtype=torch.long, device=device) for ctx_id in ctx_ids
|
|
]
|
|
ctx_attn_mask = [torch.ones_like(ids) for ids in ctx_ids]
|
|
ctx_attn_mask = [
|
|
torch.tensor(mask, dtype=torch.long, device=device) for mask in ctx_attn_mask
|
|
]
|
|
ctx_ids = torch.nn.utils.rnn.pad_sequence(
|
|
ctx_ids,
|
|
batch_first=True,
|
|
padding_value=0,
|
|
)
|
|
ctx_attn_mask = torch.nn.utils.rnn.pad_sequence(
|
|
ctx_attn_mask,
|
|
batch_first=True,
|
|
padding_value=0,
|
|
)
|
|
return {"ctx_ids": ctx_ids, "ctx_attn_mask": ctx_attn_mask}
|
|
|
|
|
|
def add_user_message(message: str, history):
|
|
if not message.strip():
|
|
return history, ""
|
|
return history + [[message, None]], ""
|
|
|
|
|
|
def generate_response(
|
|
history: list[list[str]],
|
|
system_msg: str,
|
|
context_1: str,
|
|
scaler_1: float,
|
|
context_2: str,
|
|
scaler_2: float,
|
|
context_3: str,
|
|
scaler_3: float,
|
|
bias_scaler: float,
|
|
):
|
|
global modulated_model, chat_history, ctx_tokenizer, base_tokenizer
|
|
|
|
if modulated_model is None:
|
|
history[-1][1] = "Please load a checkpoint first."
|
|
yield history
|
|
return
|
|
|
|
if not history or history[-1][0] is None:
|
|
yield history
|
|
return
|
|
|
|
try:
|
|
user_message = history[-1][0]
|
|
|
|
if system_msg.strip() and chat_history[0]["role"] == "system":
|
|
chat_history[0]["content"] = system_msg.strip()
|
|
|
|
chat_history.append({"role": "user", "content": user_message})
|
|
|
|
contexts = []
|
|
scalers = []
|
|
|
|
if context_1 and context_1.strip():
|
|
contexts.append(context_1)
|
|
scalers.append(scaler_1)
|
|
if context_2 and context_2.strip():
|
|
contexts.append(context_2)
|
|
scalers.append(scaler_2)
|
|
if context_3 and context_3.strip():
|
|
contexts.append(context_3)
|
|
scalers.append(scaler_3)
|
|
|
|
if not contexts:
|
|
contexts = [""]
|
|
scalers = [1.0]
|
|
|
|
print(f"Processing {len(contexts)} contexts with scalers: {scalers}")
|
|
print(f"Bias scaler: {bias_scaler}")
|
|
|
|
with torch.inference_mode(), torch.amp.autocast(str(device)):
|
|
ctx_inputs = process_multiple_contexts(contexts)
|
|
ctx_ids = ctx_inputs["ctx_ids"].to(device)
|
|
ctx_attn_mask = ctx_inputs["ctx_attn_mask"].to(device)
|
|
|
|
scalers_tensor = torch.tensor(scalers, dtype=torch.float32, device=device)
|
|
|
|
model_inputs = base_tokenizer.apply_chat_template(
|
|
chat_history, return_tensors="pt", add_generation_prompt=True
|
|
).to(device)
|
|
|
|
print(f"Contexts: {contexts}")
|
|
print(f"Chat history: {chat_history}")
|
|
|
|
outputs = modulated_model.generate(
|
|
ctx_ids=ctx_ids,
|
|
ctx_attn_mask=ctx_attn_mask,
|
|
n_ctx_chunks=torch.tensor([len(ctx_ids)], device=ctx_ids.device),
|
|
scalers=scalers_tensor,
|
|
bias_scaler=bias_scaler,
|
|
input_ids=model_inputs,
|
|
max_new_tokens=512,
|
|
do_sample=False,
|
|
temperature=0,
|
|
)
|
|
|
|
response = base_tokenizer.decode(
|
|
outputs[0][model_inputs.shape[1] :], skip_special_tokens=True
|
|
)
|
|
|
|
chat_history.append({"role": "assistant", "content": response})
|
|
|
|
words = response.split()
|
|
partial_response = ""
|
|
for word in words:
|
|
partial_response += word + " "
|
|
history[-1][1] = partial_response.strip()
|
|
yield history
|
|
|
|
history[-1][1] = response
|
|
yield history
|
|
|
|
except Exception as e:
|
|
import traceback
|
|
|
|
error_msg = f"❌ Error: {str(e)}"
|
|
print(f"Error generating response: {str(e)}\n\n{traceback.format_exc()}")
|
|
history[-1][1] = error_msg
|
|
yield history
|
|
|
|
|
|
def reset_chat(system_msg: str):
|
|
global chat_history
|
|
chat_history = [
|
|
{"role": "system", "content": system_msg.strip() if system_msg else ""}
|
|
]
|
|
return [[None, WARNING_MESSAGE]], "Chat history reset successfully!"
|
|
|
|
|
|
def update_context_display(num_contexts: int):
|
|
updates = []
|
|
for i in range(3):
|
|
updates.append(gr.update(visible=(i < num_contexts)))
|
|
return updates
|
|
|
|
|
|
custom_css = """
|
|
:root {
|
|
color-scheme: light;
|
|
}
|
|
|
|
.gradio-container {
|
|
font-family: 'Inter', sans-serif;
|
|
}
|
|
|
|
.chat-container {
|
|
border-radius: 10px;
|
|
border: 1px solid #e0e0e0;
|
|
}
|
|
|
|
.context-field {
|
|
background-color: #f9f9f9;
|
|
border-radius: 8px;
|
|
padding: 15px;
|
|
margin-bottom: 10px;
|
|
border: 1px solid #e0e0e0;
|
|
}
|
|
|
|
.status-box {
|
|
border-radius: 8px;
|
|
padding: 15px;
|
|
margin: 10px 0;
|
|
}
|
|
|
|
.primary-button {
|
|
background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
|
|
border: none;
|
|
border-radius: 6px;
|
|
color: white;
|
|
font-weight: 600;
|
|
}
|
|
|
|
.secondary-button {
|
|
background-color: #607d8b;
|
|
border: none;
|
|
border-radius: 6px;
|
|
color: white;
|
|
}
|
|
|
|
#chatbot {
|
|
height: 500px;
|
|
}
|
|
|
|
.instruction-text {
|
|
font-style: italic;
|
|
color: #666;
|
|
font-size: 0.9em;
|
|
margin-bottom: 10px;
|
|
}
|
|
|
|
.warning-box {
|
|
background-color: #fff3cd;
|
|
border: 1px solid #ffc107;
|
|
border-radius: 8px;
|
|
padding: 12px;
|
|
margin: 10px 0;
|
|
color: #856404;
|
|
font-size: 0.95em;
|
|
}
|
|
|
|
.warning-box strong {
|
|
color: #d97706;
|
|
}
|
|
"""
|
|
|
|
|
|
def create_demo():
|
|
with gr.Blocks(
|
|
title="Context-to-LoRA Chat Interface",
|
|
theme=gr.themes.Soft(),
|
|
css=custom_css,
|
|
) as demo:
|
|
gr.Markdown(
|
|
"""
|
|
# 🎯 Context-to-LoRA Chat Interface
|
|
|
|
Load a hypernetwork checkpoint and chat with a context-modulated language model.
|
|
Add multiple contexts with individual scaling parameters to influence the model's responses.
|
|
"""
|
|
)
|
|
|
|
gr.Markdown(
|
|
"""
|
|
### 📖 Usage Instructions
|
|
|
|
1. **Load a Checkpoint**: Select a hypernetwork checkpoint from the dropdown and click "Load Checkpoint"
|
|
2. **Configure Contexts**:
|
|
- Use the slider to choose how many contexts you want (0-3)
|
|
- Enter your context information in each text field
|
|
- Adjust the scaling sliders to control the influence of each context
|
|
3. **Set Bias Scaler**: Adjust the bias scaler to control overall model behavior
|
|
4. **Start Chatting**: Once the model is loaded, type your message and press Shift+Enter or click Send
|
|
5. **Reset**: Use the "Reset Chat" button to start a new conversation
|
|
|
|
💡 **Tip**: You can use context to provide background information or specific knowledge
|
|
that should influence the model's responses.
|
|
"""
|
|
)
|
|
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
gr.Markdown("### 📦 Load Checkpoint")
|
|
gr.Markdown(
|
|
"*Select a trained hypernetwork checkpoint to begin.*",
|
|
elem_classes="instruction-text",
|
|
)
|
|
|
|
checkpoint_dropdown = gr.Dropdown(
|
|
choices=get_available_checkpoints(),
|
|
label="Select Checkpoint",
|
|
value=None,
|
|
interactive=True,
|
|
)
|
|
|
|
load_btn = gr.Button("Load Checkpoint", variant="primary", size="lg")
|
|
|
|
status_box = gr.Textbox(
|
|
label="Status",
|
|
lines=8,
|
|
interactive=False,
|
|
elem_classes="status-box",
|
|
)
|
|
|
|
gr.Markdown("---")
|
|
|
|
gr.Markdown("### 📝 Context Configuration")
|
|
gr.Markdown(
|
|
"*Add contexts to influence the model's responses. Each context can be scaled independently.*",
|
|
elem_classes="instruction-text",
|
|
)
|
|
|
|
num_contexts = gr.Slider(
|
|
minimum=0,
|
|
maximum=3,
|
|
step=1,
|
|
value=1,
|
|
label="Number of Contexts",
|
|
interactive=True,
|
|
)
|
|
|
|
with gr.Group(visible=True) as context_1_group:
|
|
context_1 = gr.Textbox(
|
|
label="Context 1",
|
|
placeholder="Enter context information here...",
|
|
lines=4,
|
|
elem_classes="context-field",
|
|
value=DEFAULT_CONTEXT,
|
|
)
|
|
scaler_1 = gr.Slider(
|
|
minimum=-2.0,
|
|
maximum=2.0,
|
|
step=0.01,
|
|
value=1.0,
|
|
label="Context 1 Scaling",
|
|
)
|
|
|
|
with gr.Group(visible=False) as context_2_group:
|
|
context_2 = gr.Textbox(
|
|
label="Context 2",
|
|
placeholder="Enter additional context information...",
|
|
lines=4,
|
|
elem_classes="context-field",
|
|
)
|
|
scaler_2 = gr.Slider(
|
|
minimum=-2.0,
|
|
maximum=2.0,
|
|
step=0.01,
|
|
value=1.0,
|
|
label="Context 2 Scaling",
|
|
)
|
|
|
|
with gr.Group(visible=False) as context_3_group:
|
|
context_3 = gr.Textbox(
|
|
label="Context 3",
|
|
placeholder="Enter additional context information...",
|
|
lines=4,
|
|
elem_classes="context-field",
|
|
)
|
|
scaler_3 = gr.Slider(
|
|
minimum=-2.0,
|
|
maximum=2.0,
|
|
step=0.01,
|
|
value=1.0,
|
|
label="Context 3 Scaling",
|
|
)
|
|
|
|
gr.Markdown("---")
|
|
|
|
bias_scaler = gr.Slider(
|
|
minimum=-2.0,
|
|
maximum=2.0,
|
|
step=0.01,
|
|
value=1.0,
|
|
label="Bias Scaler",
|
|
info="A single scalar applied to bias parameters (independent of contexts)",
|
|
)
|
|
|
|
with gr.Column(scale=2):
|
|
gr.Markdown("### 💬 Chat Interface")
|
|
|
|
system_msg = gr.Textbox(
|
|
label="System Message",
|
|
placeholder="Enter system message (optional)...",
|
|
lines=2,
|
|
)
|
|
|
|
chatbot = gr.Chatbot(
|
|
label="Conversation",
|
|
height=500,
|
|
elem_id="chatbot",
|
|
elem_classes="chat-container",
|
|
value=[[None, WARNING_MESSAGE]],
|
|
)
|
|
|
|
with gr.Row():
|
|
msg = gr.Textbox(
|
|
label="Your Message",
|
|
placeholder="Type your message here... (Shift+Enter for new line)",
|
|
lines=2,
|
|
scale=4,
|
|
interactive=False,
|
|
)
|
|
send_btn = gr.Button(
|
|
"Send", variant="primary", scale=1, interactive=False
|
|
)
|
|
|
|
with gr.Row():
|
|
clear_btn = gr.Button("🔄 Reset Chat", variant="secondary")
|
|
reset_status = gr.Textbox(label="Reset Status", visible=False)
|
|
|
|
load_btn.click(
|
|
fn=load_checkpoint,
|
|
inputs=[checkpoint_dropdown],
|
|
outputs=[status_box, msg],
|
|
).then(fn=lambda: gr.update(interactive=True), outputs=[send_btn])
|
|
|
|
num_contexts.change(
|
|
fn=update_context_display,
|
|
inputs=[num_contexts],
|
|
outputs=[context_1_group, context_2_group, context_3_group],
|
|
)
|
|
|
|
msg.submit(
|
|
fn=add_user_message,
|
|
inputs=[msg, chatbot],
|
|
outputs=[chatbot, msg],
|
|
).then(
|
|
fn=generate_response,
|
|
inputs=[
|
|
chatbot,
|
|
system_msg,
|
|
context_1,
|
|
scaler_1,
|
|
context_2,
|
|
scaler_2,
|
|
context_3,
|
|
scaler_3,
|
|
bias_scaler,
|
|
],
|
|
outputs=[chatbot],
|
|
)
|
|
|
|
send_btn.click(
|
|
fn=add_user_message,
|
|
inputs=[msg, chatbot],
|
|
outputs=[chatbot, msg],
|
|
).then(
|
|
fn=generate_response,
|
|
inputs=[
|
|
chatbot,
|
|
system_msg,
|
|
context_1,
|
|
scaler_1,
|
|
context_2,
|
|
scaler_2,
|
|
context_3,
|
|
scaler_3,
|
|
bias_scaler,
|
|
],
|
|
outputs=[chatbot],
|
|
)
|
|
|
|
clear_btn.click(
|
|
fn=reset_chat,
|
|
inputs=[system_msg],
|
|
outputs=[chatbot, reset_status],
|
|
)
|
|
|
|
gr.Markdown(f"---\n{FOOTER.strip()}")
|
|
return demo
|
|
|
|
|
|
if __name__ == "__main__":
|
|
demo = create_demo()
|
|
demo.launch(
|
|
server_name="0.0.0.0",
|
|
server_port=7860,
|
|
share=False,
|
|
debug=True,
|
|
)
|