update inference logic

app.py

@@ -6,30 +6,57 @@ import glob
 import os
 import uuid
 from pathlib import Path
-import spaces
-import torch
-import transformers
+from huggingface_hub import InferenceClient
+from openai import OpenAI
+from huggingface_hub import get_token
 from huggingface_hub import CommitScheduler, hf_hub_download, login
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from outlines import models, generate
-from gradio import update
 
-model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)
+from prompts import detailed_genre_description_prompt, basic_prompt
+import random
 
+# TODOs
+# 1. Add a login button
+# 2. Prompt library expand
+# 3. log user if logged in
+
+
+client = OpenAI(
+    base_url="https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct/v1",
+    api_key=get_token(),
+)
+
+
+def generate_prompt():
+    if random.choice([True, False]):
+        return detailed_genre_description_prompt()
+    else:
+        return basic_prompt()
+
+
+def generate_blurb():
+    max_tokens = random.randint(100, 1000)
+    prompt = generate_prompt()
+    print(prompt)
+    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
 
-@spaces.GPU(duration=120)
-def generate_blurb(history):
-    model = models.transformers(model_id)
-    generator = generate.text(model)
-    resp = generator("Write a blurb for a book")
-    return resp
 
 # Function to log blurb and vote
 def log_blurb_and_vote(blurb, vote):
     log_entry = {"timestamp": datetime.now().isoformat(), "blurb": blurb, "vote": vote}
     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}"
 
 
@@ -38,19 +65,28 @@ tufte_theme = TufteInspired()
 
 # Create Gradio interface
 with gr.Blocks(theme=tufte_theme) as demo:
-    gr.Markdown("<h1 style='text-align: center;'>Would you read it?</h1>")
+    gr.Markdown("<h1 style='text-align: center;'>Would you read this book?</h1>")
     gr.Markdown(
-        "Click the button to generate a blurb for a made-up book, then vote on its quality."
+        """<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://example.com">this</a> Hugging Face dataset</p>"""
+        + """"""
     )
+    # gr.LoginButton(size="sm")
+    user_name = gr.Textbox(label="User Name", placeholder="Enter your name")
     with gr.Row():
-        generate_btn = gr.Button("Write a Blurb", variant="primary")
-    blurb_output = gr.Textbox(label="Generated Blurb", lines=5, interactive=False)
-    with gr.Row():
+        generate_btn = gr.Button("Create a book", variant="primary")
+    blurb_output = gr.Markdown(label="Book blurb")
+    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)
+    vote_output = gr.Textbox(label="Vote Status", interactive=False, visible=False)
+
+    def show_voting_buttons(blurb):
+        return blurb, gr.Row(visible=True)
 
-    generate_btn.click(generate_blurb, outputs=blurb_output)
+    generate_btn.click(generate_blurb, outputs=blurb_output).then(
+        show_voting_buttons, inputs=blurb_output, outputs=[blurb_output, voting_row]
+    )
     upvote_btn.click(
         lambda x: log_blurb_and_vote(x, "upvote"),
         inputs=blurb_output,
@@ -62,5 +98,6 @@ with gr.Blocks(theme=tufte_theme) as demo:
         outputs=vote_output,
     )
 
+
 if __name__ == "__main__":
     demo.launch()