app.py

import gradio as gr
import transformers
import torch
import json
from transformers import AutoTokenizer
import os
from huggingface_hub import login
import spaces

HF_TOKEN = os.getenv("HF_TOKEN")
login(HF_TOKEN)
# Load the model
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)

pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.bfloat16},
    device="cuda",
)

# Load the model configuration
with open("model_configs.json", "r") as f:
    model_configs = json.load(f)
    model_config = model_configs[model_id]

# Extract instruction
extract_input = model_config["extract_input"]
terminators = [
    tokenizer.eos_token_id,
    tokenizer.convert_tokens_to_ids("<|eot_id|>"),
]


@spaces.GPU
def generate_instruction_response():
    prompt_info = f"""### Generating user prompt using the template:

```
{extract_input}
```
"""
    yield prompt_info
    instruction = pipeline(
        extract_input,
        max_new_tokens=2048,
        eos_token_id=terminators,
        do_sample=True,
        temperature=1,
        top_p=1,
    )

    sanitized_instruction = instruction[0]["generated_text"][
        len(extract_input) :
    ].split("\n")[0]

    first_step = (
        prompt_info + f"### LLM generated instruction:\n\n{sanitized_instruction}"
    )
    yield first_step + "\n\n### Generating LLM response..."

    response_template = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{sanitized_instruction}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"""

    response = pipeline(
        response_template,
        max_new_tokens=2048,
        eos_token_id=terminators,
        do_sample=True,
        temperature=1,
        top_p=1,
    )

    assistant_response = response[0]["generated_text"][len(response_template) :]

    final_output = f"""### Template used for generating instruction:

```
{extract_input}
```

### LLM Generated Instruction:

{sanitized_instruction}

### Template used for generating response:

```
{response_template}
```

### LLM Generated Response:

{assistant_response}
"""
    yield final_output


title = "Magpie Demo"
description = """
This Gradio demo showcases the approach described in the Magpie paper. Magpie is a data synthesis pipeline that creates high-quality alignment data without relying on prompt engineering or seed questions. Instead, it generates instruction data by prompting aligned LLMs with a pre-query template.

In this demo, you can see how the model generates a user instruction and a model response, along with the templates used in the process.

You can learn more about the approach [in the paper](https://huggingface.co/papers/2406.08464).
"""
# Create the Gradio interface
iface = gr.Interface(
    fn=generate_instruction_response,
    inputs=[],
    outputs=[gr.Markdown(label="Generated Data")],
    title=title,
    description=description,
    submit_btn="Generate Instructions Response Pair",
)

# Launch the app
iface.launch(debug=True)