chore: Refactor code for improved UI layout and readability

app.py

@@ -1,72 +1,23 @@
-import hashlib
+import gradio as gr
 import json
-import os
-import random
-import uuid
 from datetime import datetime
-from pathlib import Path
-
-import gradio as gr
-from huggingface_hub import CommitScheduler, get_token, login
-from openai import OpenAI
-
-from prompts import basic_prompt, detailed_genre_description_prompt
 from theme import TufteInspired
+import uuid
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+from huggingface_hub import get_token, login
+from prompts import detailed_genre_description_prompt, basic_prompt
+import random
+import os
 
 # Ensure you're logged in to Hugging Face
-login(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",
-]
+login(get_token())
 
-
-def get_random_model():
-    return random.choice(MODELS)
-
-
-def create_client(model_id):
-    return OpenAI(
-        base_url=f"https://api-inference.huggingface.co/models/{model_id}/v1",
-        api_key=get_token(),
-    )
-
-
-# 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-preferences"
-scheduler = CommitScheduler(
-    repo_id=repo_id,
-    repo_type="dataset",
-    folder_path=dataset_folder,
-    path_in_repo="data",
-    every=1,  # Upload every minute
+client = OpenAI(
+    base_url="https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct/v1",
+    api_key=get_token(),
 )
 
-# Global dictionary to store votes
-votes = {}
-
 
 def generate_prompt():
     if random.choice([True, False]):
@@ -82,8 +33,6 @@ def get_and_store_prompt():
 
 
 def generate_blurb(prompt):
-    model_id = get_random_model()
-    client = create_client(model_id)
     max_tokens = random.randint(100, 1000)
     chat_completion = client.chat.completions.create(
         model="tgi",
@@ -96,44 +45,23 @@ def generate_blurb(prompt):
     full_text = ""
     for message in chat_completion:
         full_text += message.choices[0].delta.content
-        yield full_text, model_id
-    return full_text, model_id  # Return final result with model_id
+        yield full_text
 
 
-def generate_vote_id(user_id, blurb):
-    return hashlib.md5(f"{user_id}:{blurb}".encode()).hexdigest()
-
-
-def log_blurb_and_vote(
-    prompt,
-    blurb,
-    vote,
-    model_id,
-    user_info: gr.OAuthProfile | None,
-):
+# 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())
-    vote_id = generate_vote_id(user_id, blurb)
-
-    if vote_id in votes:
-        gr.Info("You've already voted on this blurb!")
-        return None, gr.Row.update(visible=False)
-
-    votes[vote_id] = vote
-
     log_entry = {
         "timestamp": datetime.now().isoformat(),
         "prompt": prompt,
         "blurb": blurb,
         "vote": vote,
         "user_id": user_id,
-        "model_id": model_id,
     }
-    with scheduler.lock:
-        with dataset_file.open("a") as f:
-            f.write(json.dumps(log_entry) + "\n")
-
-    gr.Info("Thank you for voting! Your feedback will be synced to the dataset.")
-    return f"Logged: {vote} by user {user_id}", gr.Row.update(visible=False)
+    with open("blurb_log.jsonl", "a") as f:
+        f.write(json.dumps(log_entry) + "\n")
+    gr.Info("Thank you for voting!")
+    return f"Logged: {vote} by user {user_id}"
 
 
 # Create custom theme
@@ -144,41 +72,34 @@ with gr.Blocks(theme=tufte_theme) as demo:
     gr.Markdown("<h1 style='text-align: center;'>Would you read this book?</h1>")
     gr.Markdown(
         """<p style='text-align: center;'>Looking for your next summer read? 
-    Would you read a book based on this LLM generated blurb? <br> Your vote will be added to <a href="https://huggingface.co/datasets/your-username/your-dataset-repo">this</a> Hugging Face dataset</p>"""
+    Would you read a book based on this LLM generated blurb? <br> Your vote will be added to <a href="https://example.com">this</a> Hugging Face dataset</p>"""
     )
 
-    with gr.Row():
-        login_btn = gr.LoginButton(size="sm")
+    # Add the login button
+    login_btn = gr.LoginButton()
+
     with gr.Row():
         generate_btn = gr.Button("Create a book", variant="primary")
 
     prompt_state = gr.State()
     blurb_output = gr.Markdown(label="Book blurb")
-    user_state = gr.State()
-    model_state = gr.State()
 
     with gr.Row(visible=False) 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=True)
+    vote_output = gr.Textbox(label="Vote Status", interactive=False, visible=False)
 
-    def generate_and_show(prompt, user_info):
-        return "Generating...", gr.Row.update(visible=False), user_info, None
+    def generate_and_show(prompt):
+        return gr.Markdown.update(value="Generating..."), gr.Row(visible=False)
 
-    def show_voting_buttons(blurb, model_id):
-        return blurb, gr.Row.update(visible=True), model_id
+    def show_voting_buttons(blurb):
+        return blurb, gr.Row(visible=True)
 
     generate_btn.click(get_and_store_prompt, outputs=prompt_state).then(
-        generate_and_show,
-        inputs=[prompt_state, login_btn],
-        outputs=[blurb_output, voting_row, user_state, model_state],
-    ).then(
-        generate_blurb, inputs=prompt_state, outputs=[blurb_output, model_state]
-    ).then(
-        show_voting_buttons,
-        inputs=[blurb_output, model_state],
-        outputs=[blurb_output, voting_row, model_state],
+        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(
@@ -187,10 +108,9 @@ with gr.Blocks(theme=tufte_theme) as demo:
             prompt_state,
             blurb_output,
             gr.Textbox(value="upvote", visible=False),
-            model_state,
-            user_state,
+            login_btn,
         ],
-        outputs=[vote_output, voting_row],
+        outputs=vote_output,
     )
     downvote_btn.click(
         log_blurb_and_vote,
@@ -198,10 +118,9 @@ with gr.Blocks(theme=tufte_theme) as demo:
             prompt_state,
             blurb_output,
             gr.Textbox(value="downvote", visible=False),
-            model_state,
-            user_state,
+            login_btn,
         ],
-        outputs=[vote_output, voting_row],
+        outputs=vote_output,
     )
 
 if __name__ == "__main__":