app.py

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import torch
from threading import Thread
import re  # For cleaning unwanted tokens

# Load model and tokenizer
model_name = "GoofyLM/gonzalez-v1"
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    torch_dtype=torch.float16
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Set pad token if missing
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# Define a custom chat template if one is not available
if tokenizer.chat_template is None:
    tokenizer.chat_template = """{% for message in messages %}
{% if message['role'] == 'system' %}<|system|>
{{ message['content'] }}
{% elif message['role'] == 'user' %}<|user|>
{{ message['content'] }}
{% elif message['role'] == 'assistant' %}<|assistant|>
{{ message['content'] }}
{% endif %}
{% endfor %}
{% if add_generation_prompt %}<|assistant|>
{% endif %}"""

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    # Build conversation messages
    messages = [{"role": "system", "content": system_message}]
    
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    
    messages.append({"role": "user", "content": message})
    
    # Format prompt using chat template
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    
    # Set up streaming
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    # Configure generation parameters
    do_sample = temperature > 0 or top_p < 1.0
    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=do_sample,
        pad_token_id=tokenizer.pad_token_id
    )
    
    # Start generation in a separate thread
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    # Stream response with filtering
    response = ""
    for token in streamer:
        response += token
        # Remove angle-bracket tags (e.g. <|msg|...>, <username: ...>, <:...:...>)
        cleaned_response = re.sub(r"<[^>]+>", "", response)
        # Remove leading "Output:" if present (case-insensitive, line start)
        cleaned_response = re.sub(r"(?i)^\s*output:\s*", "", cleaned_response)
        yield cleaned_response.strip()

# Create Gradio interface
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="", label="System message"),
        gr.Slider(1, 2048, value=72, label="Max new tokens"),
        gr.Slider(0.1, 4.0, value=0.7, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.95, label="Top-p (nucleus sampling)"),
    ],
)

if __name__ == "__main__":
    demo.launch()