Update app.py

app.py

@@ -1,76 +1,66 @@
 import gradio as gr
 import json
-import random
 from datetime import datetime
 from theme import TufteInspired
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
-import torch
+import glob
 import os
+import uuid
+from pathlib import Path
 import spaces
+import torch
+import transformers
+from huggingface_hub import CommitScheduler, hf_hub_download, login
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from outlines import models, generate
+from gradio import update
 
-HF_TOKEN = os.getenv("HF_TOKEN")
-login(HF_TOKEN)
-
-# List of models to choose from
-model_list = [
-    "meta-llama/Llama-2-7b-chat-hf",
-]
-
-# Function to load a random model
-@spaces.GPU(duration=120)  # Allowing extra time for model loading
-def load_random_model():
-    model_id = random.choice(model_list)
-    tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
-    return model_id, model, tokenizer
+model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)
 
 
-@spaces.GPU
-def generate_blurb():
-    model_id, model, tokenizer = load_random_model()
-    prompt = "Write a blurb for a made-up book:"
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=200, num_return_sequences=1)
-    
-    blurb = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return f"Model used: {model_id}\n\nBlurb: {blurb}"
+@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
-    }
+    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")
     return f"Logged: {vote}"
 
+
 # Create custom theme
 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("Click the button to generate a blurb for a made-up book using a random model, then vote on its quality.")
-    
+    gr.Markdown(
+        "Click the button to generate a blurb for a made-up book, then vote on its quality."
+    )
     with gr.Row():
         generate_btn = gr.Button("Write a Blurb", variant="primary")
-    
-    blurb_output = gr.Textbox(label="Generated Blurb", lines=8, interactive=False)
-    
+    blurb_output = gr.Textbox(label="Generated Blurb", lines=5, interactive=False)
     with gr.Row():
         upvote_btn = gr.Button("👍 would read")
         downvote_btn = gr.Button("👎 wouldn't read")
-    
     vote_output = gr.Textbox(label="Vote Status", interactive=False)
-    
+
     generate_btn.click(generate_blurb, outputs=blurb_output)
-    upvote_btn.click(lambda x: log_blurb_and_vote(x, "upvote"), inputs=blurb_output, outputs=vote_output)
-    downvote_btn.click(lambda x: log_blurb_and_vote(x, "downvote"), inputs=blurb_output, outputs=vote_output)
+    upvote_btn.click(
+        lambda x: log_blurb_and_vote(x, "upvote"),
+        inputs=blurb_output,
+        outputs=vote_output,
+    )
+    downvote_btn.click(
+        lambda x: log_blurb_and_vote(x, "downvote"),
+        inputs=blurb_output,
+        outputs=vote_output,
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()