app.py

import json
import os
import random
import uuid
from datetime import datetime
from pathlib import Path
from typing import Callable, List, Tuple

import gradio as gr
from huggingface_hub import (
    CommitScheduler,
    InferenceClient,
    get_token,
    hf_hub_download,
    login,
)
from openai import OpenAI

from prompts import basic_prompt, detailed_genre_description_prompt, very_basic_prompt
from theme import TufteInspired

HF_TOKEN = os.getenv("HF_TOKEN")

# Define available models
MODELS = [
    "meta-llama/Meta-Llama-3-70B-Instruct",
    "mistralai/Mixtral-8x7B-Instruct-v0.1",
    "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
]


# Set up dataset storage
dataset_folder = Path("dataset")
dataset_folder.mkdir(exist_ok=True)


# Function to get the latest dataset file
def get_latest_dataset_file():
    files = list(dataset_folder.glob("data_*.jsonl"))
    return max(files, key=os.path.getctime) if files else None


# Check for existing dataset and create or append to it
if latest_file := get_latest_dataset_file():
    dataset_file = latest_file
    print(f"Appending to existing dataset file: {dataset_file}")
else:
    dataset_file = dataset_folder / f"data_{uuid.uuid4()}.jsonl"
    print(f"Creating new dataset file: {dataset_file}")

# Set up CommitScheduler for dataset uploads
repo_id = (
    "davanstrien/summer-reading-preference"  # Replace with your desired dataset repo
)
scheduler = CommitScheduler(
    repo_id=repo_id,
    repo_type="dataset",
    folder_path=dataset_folder,
    path_in_repo="data",
    every=5,  # Upload every 5 minutes
)


# Function to download existing dataset files
def download_existing_dataset():
    try:
        files = hf_hub_download(
            repo_id=repo_id,
            filename="data",
            repo_type="dataset",
        )
        for file in Path(files).glob("*.jsonl"):
            dest_file = dataset_folder / file.name
            if not dest_file.exists():
                dest_file.write_bytes(file.read_bytes())
                print(f"Downloaded existing dataset file: {dest_file}")
    except Exception as e:
        print(f"Error downloading existing dataset: {e}")


# Download existing dataset files at startup
download_existing_dataset()


def get_random_model():
    global CHOSEN_MODEL
    model = random.choice(MODELS)
    CHOSEN_MODEL = model
    return model


def create_client(model_id):
    return OpenAI(
        base_url=f"https://api-inference.huggingface.co/models/{model_id}/v1",
        api_key=HF_TOKEN,
    )


# client = OpenAI(
#     base_url="https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct/v1",
#     api_key=HF_TOKEN
# )


def weighted_random_choice(choices: List[Tuple[Callable, float]]) -> Callable:
    total = sum(weight for _, weight in choices)
    r = random.uniform(0, total)
    upto = 0
    for choice, weight in choices:
        if upto + weight >= r:
            return choice
        upto += weight
    assert False, "Shouldn't get here"


def generate_prompt() -> str:
    prompt_choices = [
        (detailed_genre_description_prompt, 0.5),
        (basic_prompt, 0.3),
        (very_basic_prompt, 0.2),
    ]

    selected_prompt_func = weighted_random_choice(prompt_choices)
    return selected_prompt_func()


def get_and_store_prompt():
    prompt = generate_prompt()
    print(prompt)  # Keep this for debugging
    return prompt


def generate_blurb(prompt):
    model_id = get_random_model()
    client = create_client(model_id)
    max_tokens = random.randint(500, 1000)
    chat_completion = client.chat.completions.create(
        model="tgi",
        messages=[
            {"role": "user", "content": prompt},
        ],
        stream=True,
        max_tokens=max_tokens,
    )
    full_text = ""
    for message in chat_completion:
        full_text += message.choices[0].delta.content
        yield full_text


# Function to log blurb and vote
def log_blurb_and_vote(prompt, blurb, vote, user_info: gr.OAuthProfile | None, *args):
    user_id = user_info.username if user_info is not None else str(uuid.uuid4())
    log_entry = {
        "timestamp": datetime.now().isoformat(),
        "prompt": prompt,
        "blurb": blurb,
        "vote": vote,
        "user_id": user_id,
        "model": CHOSEN_MODEL,
    }
    with scheduler.lock:
        with dataset_file.open("a") as f:
            f.write(json.dumps(log_entry) + "\n")
    gr.Info("Thank you for voting!")
    return f"Logged: {vote} by user {user_id}", gr.Row.update(visible=False)


short_description = """Vote on book blurbs generated by large language models. Would you read the book the LLM generated? <br> Every five minutes, the dataset of votes created in this will be uploaded to the <a href="https://huggingface.co/datasets/davanstrien/summer-reading-preference">davanstrien/summer-reading-preference</a> dataset.
"""

full_description = """
Large Language Models are already strong assistants for technical tasks like coding. Increasingly they are also being used to help with tasks like copywriting. The jury is out on whether the texts produced by language models in these applications is very appealing; "let's delve into" is a common example of the clunky cliche ridden text that LLMs often produce.  <br>

However, there is also growing interest in using LLMs to help with more creative tasks. Outside of larger companies, there is a growing community of people fine-tuning LLMs for all sorts of creative tasks. Some writers want to use LLMs not as a replacement but as a companion in their writing process. <br>

One of the requirements for building models which are better able to generate responses people like is having preference data. Preference datasets come in many forms but essentially boil to a dataset which contains some kind of signal for whether people like or dislike some LLM generated text. <br>

This Space is a small experiment to see if we can generate preference data for LLM generated book blurbs. Whilst writing a blurb is very different from writing a whole book, it could be a neat experiment to see whether we can improve the ability of LLMs to generate book blurbs that people like. <br>
"""

# Create custom theme
tufte_theme = TufteInspired()
with gr.Blocks(theme=tufte_theme) as demo:
    gr.Markdown("<h1 style='text-align: center;'>Would you read this book?</h1>")
    gr.Markdown(f"""<p style='text-align: center;'>{short_description}</p>""")
    with gr.Accordion("More information", open=False):
        gr.Markdown(full_description)
    with gr.Accordion("View the data generated so far...", open=False):
        gr.Markdown("Below is the progress of the dataset so far!")
        gr.HTML("""<iframe
    src="https://huggingface.co/datasets/davanstrien/summer-reading-preference/embed/viewer"
    frameborder="0"
    width="100%"
  height="560px"
></iframe>""")
    with gr.Row():
        login_btn = gr.LoginButton(size="sm")
        gr.Markdown(
            "Login with your Hugging Face account to assign your Hub username to your votes. This will allow you to extract your preferences from the dataset generated by this Space! If you don't have a Hugging Face account, you can still vote but your votes will be anonymous."
        )
    with gr.Row():
        generate_btn = gr.Button("Create a book", variant="primary")

    prompt_state = gr.State()
    blurb_output = gr.Markdown(label="Book blurb")

    with gr.Row(visible=True) as voting_row:
        upvote_btn = gr.Button("👍 would read")
        downvote_btn = gr.Button("👎 wouldn't read")

    vote_output = gr.Textbox(label="Vote Status", interactive=False, visible=False)

    def generate_and_show(prompt):
        return "Generating...", gr.Row.update(visible=False)

    def show_voting_buttons(blurb):
        return blurb, gr.Row.update(visible=True)

    generate_btn.click(get_and_store_prompt, outputs=prompt_state).then(
        generate_and_show, inputs=prompt_state, outputs=[blurb_output, voting_row]
    ).then(generate_blurb, inputs=prompt_state, outputs=blurb_output).then(
        show_voting_buttons, inputs=blurb_output, outputs=[blurb_output, voting_row]
    )

    upvote_btn.click(
        log_blurb_and_vote,
        inputs=[
            prompt_state,
            blurb_output,
            gr.Textbox(value="upvote", visible=False),
            login_btn,
        ],
        outputs=[vote_output, voting_row],
    )
    downvote_btn.click(
        log_blurb_and_vote,
        inputs=[
            prompt_state,
            blurb_output,
            gr.Textbox(value="downvote", visible=False),
            login_btn,
        ],
        outputs=[vote_output, voting_row],
    )


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