fix

app.py

@@ -16,22 +16,36 @@ from theme import TufteInspired
 # 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",
+]
 
 
-client = OpenAI(
-    base_url="https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-70B-Instruct/v1",
-    api_key=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
@@ -53,18 +67,23 @@ scheduler = CommitScheduler(
 # Global dictionary to store votes
 votes = {}
 
+
 def generate_prompt():
     if random.choice([True, False]):
         return detailed_genre_description_prompt()
     else:
         return basic_prompt()
 
+
 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(100, 1000)
     chat_completion = client.chat.completions.create(
         model="tgi",
@@ -77,12 +96,17 @@ def generate_blurb(prompt):
     full_text = ""
     for message in chat_completion:
         full_text += message.choices[0].delta.content
-        yield full_text
+        yield full_text, model_id
+    return full_text, model_id  # Return final result with model_id
+
 
 def generate_vote_id(user_id, blurb):
     return hashlib.md5(f"{user_id}:{blurb}".encode()).hexdigest()
 
-def log_blurb_and_vote(prompt, blurb, vote, user_info: gr.OAuthProfile | None, *args):
+
+def log_blurb_and_vote(
+    prompt, blurb, vote, user_info: gr.OAuthProfile | None, model_id, *args
+):
     user_id = user_info.username if user_info is not None else str(uuid.uuid4())
     vote_id = generate_vote_id(user_id, blurb)
 
@@ -98,14 +122,16 @@ def log_blurb_and_vote(prompt, blurb, vote, user_info: gr.OAuthProfile | None, *
         "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)
 
+
 # Create custom theme
 tufte_theme = TufteInspired()
 
@@ -125,6 +151,7 @@ with gr.Blocks(theme=tufte_theme) as demo:
     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")
@@ -133,20 +160,21 @@ with gr.Blocks(theme=tufte_theme) as demo:
     vote_output = gr.Textbox(label="Vote Status", interactive=False, visible=True)
 
     def generate_and_show(prompt, user_info):
-        # Optionally clear votes for the previous blurb if needed
-        # global votes
-        # votes = {k: v for k, v in votes.items() if not k.endswith(hash(previous_blurb))}
-        return "Generating...", gr.Row.update(visible=False), user_info
+        return "Generating...", gr.Row.update(visible=False), user_info, None
 
-    def show_voting_buttons(blurb):
-        return blurb, gr.Row.update(visible=True)
+    def show_voting_buttons(blurb, model_id):
+        return blurb, gr.Row.update(visible=True), model_id
 
     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],
-    ).then(generate_blurb, inputs=prompt_state, outputs=blurb_output).then(
-        show_voting_buttons, inputs=blurb_output, outputs=[blurb_output, voting_row]
+        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],
     )
 
     upvote_btn.click(
@@ -156,6 +184,7 @@ with gr.Blocks(theme=tufte_theme) as demo:
             blurb_output,
             gr.Textbox(value="upvote", visible=False),
             user_state,
+            model_state,
         ],
         outputs=[vote_output, voting_row],
     )
@@ -166,6 +195,7 @@ with gr.Blocks(theme=tufte_theme) as demo:
             blurb_output,
             gr.Textbox(value="downvote", visible=False),
             user_state,
+            model_state,
         ],
         outputs=[vote_output, voting_row],
     )