Update app.py

app.py

@@ -1,34 +1,42 @@
-
 import gradio as gr
 import json
+import random
 from datetime import datetime
 from theme import TufteInspired
-from transformers import AutoTokenizer
-from transformers import pipeline
-import torch
-
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
+import torch
 import os
+import spaces
+
 HF_TOKEN = os.getenv("HF_TOKEN")
 login(HF_TOKEN)
 
-# Load the model
-model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_id, add_special_tokens=True)
-
-pipeline = pipeline(
-    "text-generation",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device="cuda",
-)
+# 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
 
 
-# Placeholder LLM function
+@spaces.GPU
 def generate_blurb():
-    # This is where you'd call your LLM model
-    return pipeline("Write a blurb for a made-up book")[0]["generated_text"]
+    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}"
 
 # Function to log blurb and vote
 def log_blurb_and_vote(blurb, vote):
@@ -47,12 +55,12 @@ 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, then vote on its quality.")
+    gr.Markdown("Click the button to generate a blurb for a made-up book using a random model, 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=5, interactive=False)
+    blurb_output = gr.Textbox(label="Generated Blurb", lines=8, interactive=False)
     
     with gr.Row():
         upvote_btn = gr.Button("👍 would read")