Add more improvements

- Variable generation
- Ratings
- Dockerfile
- Hugging Face meta

Dockerfile

@@ -0,0 +1,33 @@
+# Build Docker image for deployment.
+
+FROM python:3.10.14-bullseye
+
+RUN apt-get update && \
+  apt-get install -y \
+  # General dependencies
+  locales \
+  locales-all && \
+  # Clean local repository of package files since they won't be needed anymore.
+  # Make sure this line is called after all apt-get update/install commands have
+  # run.
+  apt-get clean && \
+  # Also delete the index files which we also don't need anymore.
+  rm -rf /var/lib/apt/lists/*
+
+ENV LC_ALL en_US.UTF-8
+ENV LANG en_US.UTF-8
+ENV LANGUAGE en_US.UTF-8
+
+# Install dependencies
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+# Create non-root user
+RUN groupadd -g 900 mesop && useradd -u 900 -s /bin/bash -g mesop mesop
+USER mesop
+
+# Add app code here
+COPY . /srv/mesop-prompt-tuner
+WORKDIR /srv/mesop-prompt-tuner
+
+CMD ["gunicorn", "--bind", "0.0.0.0:8080", "--timeout", "60", "main:me"]

README.md

@@ -1,3 +1,12 @@
+---
+title: Mesop Prompt Tuner
+emoji: 🎸
+colorFrom: red
+colorTo: orange
+sdk: docker
+app_port: 8080
+---
+
 # Mesop Prompt Tuner
 
 Prompt tuner UI built using [Mesop](https://google.github.io/mesop/). This is a

components/table.py

@@ -1,10 +1,12 @@
+from typing import Callable
+
 import mesop as me
 
 _NUM_REQUIRED_ROWS = 3
 
 
 @me.component
-def prompt_eval_table(prompt):
+def prompt_eval_table(prompt, on_select_rating: Callable | None = None):
   """Creates a grid table for displaying and comparing different prompt version runs."""
   # Add a row for each variable
   num_vars = len(prompt.variables)
@@ -17,15 +19,17 @@ def prompt_eval_table(prompt):
       if num_vars
       else "1fr 20fr 1fr",
       margin=me.Margin.all(15),
+      overflow_x="scroll",
     )
   ):
     # Render first row. This row only displays the Prompt version.
     for i in range(table_size):
       with me.box(
         style=me.Style(
-          background="#fff",
+          background="#FFF",
           border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
           color="#000",
+          font_size=15,
           font_weight="bold",
           padding=me.Padding.all(10),
         )
@@ -38,18 +42,27 @@ def prompt_eval_table(prompt):
     # Render second row. This row only displays the headers of the table:
     # variable names, model response, avg rating.
     header_row = [""] + prompt.variables + ["Model response"] + [""]
-    for header_text in header_row:
+    for i, header_text in enumerate(header_row):
       with me.box(
         style=me.Style(
           background="#FFF",
           border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-          color="#0063FF" if header_text and header_text != "Model response" else "#333",
+          color="#0063FF" if header_text and header_text != "Model response" else "#444",
+          font_size=13,
+          font_weight="bold",
           padding=me.Padding.all(10),
         )
       ):
         # Handle the variable header case.
         if header_text and header_text != "Model response":
           me.text("{{" + header_text + "}}")
+        elif i == table_size - 1:
+          avg_rating = _calculate_avg_rating_from_prompt(prompt)
+          if avg_rating is not None:
+            with me.tooltip(message="Average rating"):
+              me.text(f"{avg_rating:.2f}", style=me.Style(text_align="center"))
+          else:
+            me.text("")
         else:
           me.text(header_text)
 
@@ -60,16 +73,59 @@ def prompt_eval_table(prompt):
         + [example["variables"][v] for v in prompt.variables]
         + [example["output"], example.get("rating", "")]
       )
-      for col_index, row in enumerate(content_row):
-        with me.box(
-          style=me.Style(
-            background="#fff",
-            border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
-            color="#000",
-            padding=me.Padding.all(10),
-          )
-        ):
-          if col_index == 0 or not row:
-            me.text(row)
-          else:
-            me.markdown(row)
+      for col_index, content in enumerate(content_row):
+        if col_index == len(content_row) - 1:
+          with me.box(
+            style=me.Style(
+              background="#FFF",
+              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
+              color="#000",
+              padding=me.Padding.all(10),
+            )
+          ):
+            me.select(
+              value=content,
+              options=[
+                me.SelectOption(label="1", value="1"),
+                me.SelectOption(label="2", value="2"),
+                me.SelectOption(label="3", value="3"),
+                me.SelectOption(label="4", value="4"),
+                me.SelectOption(label="5", value="5"),
+              ],
+              on_selection_change=on_select_rating,
+              key=f"rating_{prompt.version}_{row_index}",
+              style=me.Style(width=60),
+            )
+        elif col_index == 0 or not content:
+          with me.box(
+            style=me.Style(
+              background="#FFF",
+              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
+              color="#000",
+              font_size=14,
+              padding=me.Padding.all(10),
+              text_align="center",
+            )
+          ):
+            me.text(content)
+        else:
+          with me.box(
+            style=me.Style(
+              background="#FFF",
+              border=me.Border.all(me.BorderSide(width=1, style="solid", color="#DEE2E6")),
+              color="#000",
+              font_size=14,
+              padding=me.Padding.all(10),
+              max_height=300,
+              min_width=300,
+              overflow_y="scroll",
+            )
+          ):
+            me.markdown(content)
+
+
+def _calculate_avg_rating_from_prompt(prompt) -> float | None:
+  ratings = [int(response["rating"]) for response in prompt.responses if response.get("rating")]
+  if ratings:
+    return sum(ratings) / float(len(ratings))
+  return None

llm.py

@@ -1,3 +1,4 @@
+import json
 import os
 
 import google.generativeai as genai
@@ -19,6 +20,23 @@ For custom user input, you can leave placeholder variables. For example, if you
 variable named EMAIL, it would like {{{{EMAIL}}}} in the resulting prompt.
 """.strip()
 
+_GENERATE_VARIABLES_PROMPT = """
+Your job is to generate data for the given placeholders: {placeholders}.
+
+The generated data should reflect the name of the placeholder.
+
+Render the output as JSON.
+
+Here is an example output for these placeholders: STORY, FEEDBACK_TYPE
+
+{{
+  "STORY": "Generated data for a story",
+  "FEEDBACK_TYPE": "Type of feedback to provide on the story"
+}}
+
+Again, please generate outputs for these placeholders: {placeholders}
+""".strip()
+
 
 def _make_model(model_name: str, temperature: float) -> genai.GenerativeModel:
   return genai.GenerativeModel(
@@ -39,10 +57,17 @@ def generate_prompt(task_description: str, model_name: str, temperature: float)
 
 
 def generate_variables(
-  prompt: str, variables: dict[str, str], model_name: str, temperature: float
+  prompt: str, variable_names: list[str], model_name: str, temperature: float
 ) -> dict[str, str]:
-  # model = _make_model(model_name, temperature)
-  pass
+  model = _make_model(model_name, temperature)
+  output = (
+    model.generate_content(
+      _GENERATE_VARIABLES_PROMPT.format(placeholders=", ".join(variable_names))
+    )
+    .text.removeprefix("```json")
+    .removesuffix("```")
+  )
+  return json.loads(output)
 
 
 def run_prompt(prompt_with_variables: str, model_name: str, temperature: float) -> str:

main.py

@@ -271,7 +271,7 @@ def app():
       with me.box(style=me.Style(grid_column="1 / -2")):
         prompt = _find_prompt(state.prompts, state.version)
         if prompt:
-          mex.prompt_eval_table(prompt)
+          mex.prompt_eval_table(prompt, on_select_rating=on_select_rating)
 
     with mex.icon_sidebar():
       if state.mode == "Prompt":
@@ -439,15 +439,15 @@ def on_click_generate_prompt(e: me.ClickEvent):
 
 
 def on_click_generate_variables(e: me.ClickEvent):
-  """Generates values for the given empty variables.
-
-  TODO: Implement this logic.
-  """
+  """Generates values for the given empty variables."""
   state = me.state(State)
   variable_names = set(_parse_variables(state.prompt))
+  generated_variables = llm.generate_variables(
+    state.prompt, variable_names, state.model, state.model_temperature
+  )
   for name, value in state.prompt_variables.items():
-    if name in variable_names and not value:
-      state.prompt_variables[name] = "Generate variable " + name
+    if name in variable_names and name in generated_variables:
+      state.prompt_variables[name] = generated_variables[name]
 
 
 def on_click_mode_toggle(e: me.ClickEvent):
@@ -456,6 +456,13 @@ def on_click_mode_toggle(e: me.ClickEvent):
   state.mode = "Eval" if state.mode == "Prompt" else "Prompt"
 
 
+def on_select_rating(e: me.SelectSelectionChangeEvent):
+  state = me.state(State)
+  _, prompt_version, response_index = e.key.split("_")
+  prompt = _find_prompt(state.prompts, int(prompt_version))
+  prompt.responses[int(response_index)]["rating"] = e.value
+
+
 # Generic event handlers