app.py

import gradio as gr
import torch
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, PeftConfig
from huggingface_hub import hf_hub_download, login
import os
import json

# Retrieve the token from the environment variable
hf_token = os.environ.get("HF_TOKEN")
if hf_token:
    login(token=hf_token)
    print("Successfully logged in to Hugging Face Hub!")
else:
    print("HF_TOKEN not found in environment variables. Cannot authenticate.")

DEFAULT_FORM_QUESTIONS = [
    [
        "SAMPLE",
        "Does the algorithmic undecidability in the Halting Problem reflect metaphysical truths about reality and real limits of physical computation as suggested by Turing, or does it instead reveal nothing more than a boundary of the rules of language and the way we use language as suggested by Wittgenstein's rule-following paradox?",
        "A synthetic response might suggest that the Halting Problem illuminates limits inherent in algorithmic processes, while Wittgenstein's insights highlight the role of language and social context in shaping our understanding of those limits - both offering valuable perspectives that contribute to a nuanced understanding of computation, knowledge, and our relationship to the world."
    ],
    [
        "TRAINING",
        "Why does Spinoza consider doubt a signal of an inadequate idea?",
        "Because doubt arises when two conflicting ideas about the same object are held, showing that neither is clear and distinct."
    ],
    [
        "TRAINING",
        "How does Dewey describe the interplay between analysis and synthesis during judgment?",
        "Analysis emphasizes significant traits while synthesis places them in an inclusive context; each perfects the other in a continuous spiral.",
    
    ],
    ["NEW", "Other", "", ""]
]

def generate_output(base_model_name, adapter_name, question, max_tokens, temp):

    if adapter_name:
        tokenizer = AutoTokenizer.from_pretrained(base_model_name)
        base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
        model = PeftModel.from_pretrained(base_model, adapter_name)
        pipe = pipeline(
            task="text-generation",
            model=model,
            device_map="auto",
            torch_dtype=torch.bfloat16,
            tokenizer=tokenizer
        ) 
    else:
        pipe = pipeline(
            task="text-generation",
            model=base_model_name,
            device_map="auto",
            torch_dtype=torch.bfloat16 
        )

    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant who answers questions."}]
        },
        {
            "role": "user",
            "content": [{"type": "text", "text": question}]
        }
    ]

    chat_template = """
        {% for message in messages %}
            {% if message['role'] == 'system' %}
                {{ message['content'][0]['text'] }}

            {% elif message['role'] == 'user' %}
                User: {{ message['content'][0]['text'] }}

            {% elif message['role'] == 'assistant' %}
                Assistant: {{ message['content'][0]['text'] }}

            {% endif %}
        {% endfor %}Assistant:"""

    # Apply the chat template to format the input for the model.
    # `tokenize=False` is used because the pipeline will handle tokenization.
    # `add_generation_prompt=True` adds the necessary prompt for the model to generate a response.
    prompt = pipe.tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True,
        chat_template=chat_template
    )

    # Pass the formatted prompt to the pipeline.
    # max_new_tokens limits the length of the generated answer.
    outputs = pipe(prompt, max_new_tokens=max_tokens, temperature=temp)

    # Extract the generated text. The output is a list of dictionaries.
    # The generated text is typically found in 'generated_text' and needs to be cleaned to remove the input prompt.
    generated_text = outputs[0]['generated_text'][len(prompt):].strip()

    # generated text may include overflow of extra generated chat if first response less than max_new_tokens limit
    generated_text_first_line = generated_text.splitlines()[0]
    
    return generated_text_first_line

def run_challenge(
        base_model,
        adapter,
        custom_adapter,
        prompt_question,
        custom_question,
        max_tokens,
        temp
    ):

    if (base_model and (adapter or custom_adapter) and (prompt_question or custom_question) and max_tokens and temp):

        if prompt_question == "Other":
            prompt_question = custom_question
    
        if adapter == "Other":
            adapter = custom_adapter
    
        # get base model output
        output_base_model = generate_output(base_model, "", prompt_question, max_tokens, temp)
        #output_base_model = "TEMP OUTAGE"
        print(output_base_model)
        
        # get adapter output
        output_adapter = generate_output(base_model, adapter, prompt_question, max_tokens, temp)
        #output_adapter = "TEMP OUTAGE"
        print(output_adapter)
    
        return output_base_model, output_adapter

    else:
        return "Incomplete form values", "Incomplete form values"

def handle_question_choice(choice):
    if choice == "Other":
        return gr.update(label="Enter new question", visible=True, scale=1) # Show the textbox
    else:
        return gr.update(visible=False) # Hide the textbox

def handle_question_training(choice):
    for sublist in DEFAULT_FORM_QUESTIONS:  # Outer loop: iterates through each sublist
        if (sublist[0] == "TRAINING") and (sublist[1] == choice):
            return gr.update(value="**Training Answer**: " + sublist[2], visible=True)
        elif (sublist[0] == "SAMPLE") and (sublist[1] == choice):
            return gr.update(value="**Sample Question Information (not used in training)**: " + sublist[2], visible=True)
    return gr.update(value="", visible=False)

def handle_adapter_choice(choice):
    if choice == "Other":
        return gr.update(visible=True) # Show the textbox
    else:
        return gr.update(visible=False) # Hide the textbox

def get_base_model(adapter_repo):
    """
    Get base model for adapter via hub
    """
    try:
        # Download adapter config from hub
        config_path = hf_hub_download(repo_id=adapter_repo, filename="adapter_config.json")
        
        with open(config_path, 'r') as f:
            adapter_config = json.load(f)
        
        base_model = adapter_config.get('base_model_name_or_path', '')
        
        if base_model:
            return base_model
        else:
            print(f"❌ Base model not found")
            return "Base model not found"
        
    except Exception as e:
        print(f"Hub base model check failed: {e}")
        return "Base model not found"

def handle_adapter_radio(choice):
    if choice != "Other":
        print("handle_adapter_radio")
        print(choice)
        model = get_base_model(choice)
        return gr.update(value=model)

with gr.Blocks() as demo:
    
    # Add a title and description for the app.
    gr.Markdown("# Transformers Text-Generation Pipeline Q&A: Adapter vs. Base Model")
    gr.Markdown("### A tool for experimenting with PEFT (Parameter-Efficient Fine-tuning) and LoRA (Low-Rank Adaptation) adapter performance.")

    with gr.Row():
        
        with gr.Column(scale=1):

            with gr.Row():
            
                prompt_adapter = gr.Radio(
                    [
                        "joshause/llama-3.2-1b-dewey-how-we-think-adapter",
                        "joshause/gemma-3-1b-pt-spinoza-treatise-emendation-intellect",
                        "Other"
                    ],
                    label="Select an adapter or enter a new adapter",
                    scale=1
                )
            with gr.Group():
                custom_adapter = gr.Textbox(label="Enter an adapter", scale=1, visible=False)
            
                submit_adapter_button = gr.Button(
                    "Get Base Model",
                    scale=1,
                    variant="secondary",
                    visible=False
                )
            with gr.Row():
                prompt_base_model = gr.Textbox(
                    label="Base model (auto-filled)",
                    max_lines=1,
                    interactive=False,
                    scale=1,
                    value=""
                )
   
            with gr.Row():  
                questions = [[sublist[0]+" QUESTION: "+sublist[1], sublist[1]] for sublist in DEFAULT_FORM_QUESTIONS]
                prompt_question = gr.Radio(questions, label="Select a question or enter a new question", scale=1)
            with gr.Group():    
                training_answer = gr.Markdown("", padding=True, visible=False)
                custom_question = gr.Textbox(label="Enter new question", visible=False, scale=1)
            
            with gr.Row():

                challenge_submit_button = gr.Button(
                    "Run Comparison",
                    scale=1,
                    variant="primary"
                )

        with gr.Column(scale=1):
            with gr.Row():
                output_max_tokens = gr.Slider(
                    1, 512, step=1, value=128, label="Max new tokens", info="max_new_tokens in text-generation pipeline limits the length of the generated answer"
                )
            with gr.Row():
                output_temp = gr.Slider(
                    0.1, 0.9, step=0.1, value=0.7, label="Temperature", info="temperature in text-generation pipeline parameter controls the randomness of the generated text"
                )
            with gr.Row():
                output_adapter = gr.Textbox(
                    label="Adapter output",
                    show_label=True,
                    lines=4, 
                    interactive=False,
                    show_copy_button=True,
                )
            with gr.Row():
                output_base_model = gr.Textbox(
                    label="Base model output",
                    show_label=True,
                    lines=4, 
                    interactive=False,
                    show_copy_button=True,
                )

    prompt_adapter.change(
        fn=handle_adapter_choice,
        inputs=prompt_adapter,
        outputs=custom_adapter
    ).then(
        fn=handle_adapter_choice,
        inputs=prompt_adapter,
        outputs=submit_adapter_button
    ).then(
        fn=handle_adapter_radio,
        inputs=prompt_adapter,
        outputs=prompt_base_model
    )

    prompt_question.change(
        fn=handle_question_choice,
        inputs=prompt_question,
        outputs=custom_question
    ).then(
        fn=handle_question_training,
        inputs=prompt_question,
        outputs=training_answer
    )
    
    submit_adapter_button.click(
        fn=get_base_model,
        inputs=custom_adapter,
        outputs=prompt_base_model
    )
    
    challenge_submit_button.click(
        fn=run_challenge,
        inputs=[
            prompt_base_model,
            prompt_adapter,
            custom_adapter,
            prompt_question,
            custom_question,
            output_max_tokens,
            output_temp
        ],
        outputs=[
            output_base_model,
            output_adapter
        ] 
    )
    
if __name__ == "__main__":
    demo.launch()