app.py

import spaces
import gradio as gr
import torch
from transformers import pipeline, BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "marcelbinz/Llama-3.1-Minitaur-8B"

bnb_4bit_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    device_map="auto",
    attn_implementation="flash_attention_2",
    quantization_config=bnb_4bit_config,
)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
)

@spaces.GPU
def infer(prompt):
    return pipe(prompt, max_new_tokens=1, do_sample=True, temperature=1.0, return_full_text=False)[0]["generated_text"]

default_experiment = """Test
1
2
3
"""

with gr.Blocks(
    fill_width=True,
    css="""
    #prompt-box textarea {height:200px}
    #answer-box textarea {height:320px}
    #info-box {margin-bottom: 1rem}      /* a little spacing */
    """
) as demo:

    # ---------- NEW: info banner ----------
    gr.Markdown(
        """
        ### How to prompt:
        - We did not employ a particular prompt template – just phrase everything in natural language.  
        - Human choices are encapsulated by "<<" and ">>" tokens.
        - Most experiments in the training data are framed in terms of button presses. If possible, it is recommended to use that style.
        - You can find examples in the Supporting Information of our paper.
        """,
        elem_id="info-box",
    )

    # (optional) add a logo or hero image
    gr.Image(
        value="https://marcelbinz.github.io/imgs/centaur.png",
        show_label=False,
        height=100,
        elem_classes="mx-auto",          # centres the image
    )
    # --------------------------------------

    with gr.Row(equal_height=True):
        inp = gr.Textbox(
            label="Prompt", elem_id="prompt-box",
            lines=12, max_lines=12, scale=3, value=default_experiment
        )
        outp = gr.Textbox(
            label="Response", elem_id="answer-box",
            lines=1, interactive=False, scale=3
        )

    run = gr.Button("Run")
    run.click(infer, inp, outp)

demo.queue().launch()