davanstrien/magpie

1import spaces2import glob3import json4import os5import uuid6from datetime import datetime7from pathlib import Path8import gradio as gr9 10import torch11import transformers12from huggingface_hub import CommitScheduler, hf_hub_download, login13from transformers import AutoTokenizer14 15print(f"Is CUDA available: {torch.cuda.is_available()}")16 17print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")18 19HF_TOKEN = os.getenv("HF_TOKEN")20login(HF_TOKEN)21 22# Load the model23model_id = "meta-llama/Meta-Llama-3-8B-Instruct"24tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)25 26pipeline = transformers.pipeline(27    "text-generation",28    model=model_id,29    model_kwargs={"torch_dtype": torch.bfloat16},30    device="cuda",31)32 33# Load the model configuration34with open("model_configs.json", "r") as f:35    model_configs = json.load(f)36    model_config = model_configs[model_id]37 38# Extract instruction39extract_input = model_config["extract_input"]40terminators = [41    tokenizer.eos_token_id,42    tokenizer.convert_tokens_to_ids("<|eot_id|>"),43]44 45# Set up dataset storage46dataset_folder = Path("dataset")47dataset_folder.mkdir(exist_ok=True)48 49 50# Function to get the latest dataset file51def get_latest_dataset_file():52    if files := glob.glob(str(dataset_folder / "data_*.jsonl")):53        return max(files, key=os.path.getctime)54    return None55 56 57# Check for existing dataset and create or append to it58if latest_file := get_latest_dataset_file():59    dataset_file = Path(latest_file)60    print(f"Appending to existing dataset file: {dataset_file}")61else:62    dataset_file = dataset_folder / f"data_{uuid.uuid4()}.jsonl"63    print(f"Creating new dataset file: {dataset_file}")64 65# Set up CommitScheduler for dataset uploads66repo_id = "davanstrien/magpie-preference"  # Replace with your desired dataset repo67scheduler = CommitScheduler(68    repo_id=repo_id,69    repo_type="dataset",70    folder_path=dataset_folder,71    path_in_repo="data",72    every=5,  # Upload every 5 minutes73)74 75 76# Function to download existing dataset files77def download_existing_dataset():78    try:79        files = hf_hub_download(80            repo_id=repo_id, filename="data", repo_type="dataset", recursive=True81        )82        for file in glob.glob(os.path.join(files, "*.jsonl")):83            dest_file = dataset_folder / os.path.basename(file)84            if not dest_file.exists():85                dest_file.write_bytes(Path(file).read_bytes())86                print(f"Downloaded existing dataset file: {dest_file}")87    except Exception as e:88        print(f"Error downloading existing dataset: {e}")89 90 91# Download existing dataset files at startup92download_existing_dataset()93 94 95# Function to generate a session ID96def generate_session_id():97    return str(uuid.uuid4())98 99 100# Function to save feedback and generated data101def save_data(generated_input, generated_response, vote, session_id):102    data = {103        "timestamp": datetime.now().isoformat(),104        "prompt": generated_input,105        "completion": generated_response,106        "label": vote,107        "session_id": session_id,108    }109    with scheduler.lock:110        with dataset_file.open("a") as f:111            f.write(json.dumps(data) + "\n")112    return "Data saved and will be uploaded to the dataset repository."113 114 115@spaces.GPU116def generate_instruction_response():117    prompt_info = f"""### Generating user prompt using the template:118 119```120{extract_input}121```122"""123    yield (124        prompt_info,125        "",126        "",127        gr.update(interactive=False),128        gr.update(interactive=False),129        "",130        gr.update(interactive=False),131    )132    instruction = pipeline(133        extract_input,134        max_new_tokens=2048,135        eos_token_id=terminators,136        do_sample=True,137        temperature=1,138        top_p=1,139    )140 141    sanitized_instruction = instruction[0]["generated_text"][142        len(extract_input) :143    ].split("\n")[0]144 145    first_step = (146        f"{prompt_info}### LLM generated instruction:\n\n{sanitized_instruction}"147    )148    yield (149        first_step + "\n\n### Generating LLM response...",150        sanitized_instruction,151        "",152        gr.update(interactive=False),153        gr.update(interactive=False),154        "",155        gr.update(interactive=False),156    )157 158    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"""159 160    response = pipeline(161        response_template,162        max_new_tokens=2048,163        eos_token_id=terminators,164        do_sample=True,165        temperature=1,166        top_p=1,167    )168 169    assistant_response = response[0]["generated_text"][len(response_template) :]170 171    final_output = f"""### Template used for generating instruction:172 173```174{extract_input}175```176 177### LLM Generated Instruction:178 179{sanitized_instruction}180 181### LLM Generated Response:182 183{assistant_response}184"""185    yield (186        final_output,187        sanitized_instruction,188        assistant_response,189        gr.update(interactive=True),190        gr.update(interactive=True),191        "",192        gr.update(interactive=True),193    )194 195 196title = """197<h1 style="text-align:center">&#x1F426; Magpie Preference</h1>198"""199 200description = """201This demo showcases **[Magpie](https://magpie-align.github.io/)**, an innovative approach to generating high-quality data by prompting aligned LLMs with their pre-query templates. Unlike many existing synthetic data generation methods, Magpie doesn't rely on prompt engineering or seed questions for generating synthetic data. Instead, it uses the prompt template of an aligned LLM to generate both the user query and an LLM response.202 203<img src="https://magpie-align.github.io/images/pipeline.png" alt="Magpie Pipeline" width="50%" align="center" />204 205*Image Source: [Magpie project page](https://magpie-align.github.io/)*206 207 208As well as providing a demo for the Magpie generations, this Space also allows you to submit a preference rating for the generated data, contributing to a crowdsourced preference dataset!209 210## 🚀 How it works211 2121. **📝 Instruction Generation:** The model generates a user instruction.2132. **💬 Response Generation:** The model generates a response to this instruction.2143. **👍👎 User Feedback (optional):** Rate the quality of the generated content and contribute to a crowdsourced preference dataset for synthetic dataset. 215 216🔗 Find the crowd-generated dataset at [davanstrien/magpie-preference](https://huggingface.co/datasets/davanstrien/magpie-preference). It's updated every 5 minutes! You can also see a preview of the dataset below!217 218📚 Learn more about Magpie in the [paper](https://huggingface.co/papers/2406.08464).219 220> **Note:** A random session ID groups your feedback. No personal information is collected.221"""222 223# Create the Gradio interface224with gr.Blocks() as iface:225    gr.HTML(title)226    gr.Markdown(description)227 228    # Add a state variable to store the session ID229    session_id = gr.State(generate_session_id)230 231    generated_input = gr.State("")232    generated_response = gr.State("")233 234    generate_btn = gr.Button("🚀 Generate Instructions Response Pair")235 236    output = gr.Markdown(label="Generated Data")237 238    with gr.Row():239        gr.Markdown("*Vote on the quality of the generated data*")240    with gr.Row():241        thumbs_down = gr.Button("👎 Thumbs Down", interactive=False)242        thumbs_up = gr.Button("👍 Thumbs Up", interactive=False)243 244    feedback_output = gr.Markdown(label="Feedback Status")245 246    def vote_and_submit(vote, input_text, response_text, session_id):247        if input_text and response_text:248            feedback = save_data(249                input_text, response_text, vote == "👍 Thumbs Up", session_id250            )251            return (252                feedback,253                gr.update(interactive=False),254                gr.update(interactive=False),255                gr.update(interactive=True),256            )257        else:258            return (259                "Please generate data before submitting feedback.",260                gr.update(interactive=True),261                gr.update(interactive=True),262                gr.update(interactive=True),263            )264 265    generate_btn.click(266        generate_instruction_response,267        inputs=[],268        outputs=[269            output,270            generated_input,271            generated_response,272            thumbs_up,273            thumbs_down,274            feedback_output,275            generate_btn,276        ],277    )278    thumbs_up.click(279        vote_and_submit,280        inputs=[281            gr.State("👍 Thumbs Up"),282            generated_input,283            generated_response,284            session_id,285        ],286        outputs=[feedback_output, thumbs_up, thumbs_down, generate_btn],287    )288    thumbs_down.click(289        vote_and_submit,290        inputs=[291            gr.State("👎 Thumbs Down"),292            generated_input,293            generated_response,294            session_id,295        ],296        outputs=[feedback_output, thumbs_up, thumbs_down, generate_btn],297    )298    gr.Markdown("### Generated Dataset")299    gr.HTML("""<iframe300    src="https://huggingface.co/datasets/davanstrien/magpie-preference/embed/viewer"301    frameborder="0"302    width="100%"303  height="560px"304></iframe>""")305 306# Launch the app307iface.launch(debug=True)308