refactoring

app.py

@@ -5,7 +5,8 @@ import os
 import random
 import json
 import re
-from typing import List
+from typing import List, Tuple
+import PIL
 
 import gradio as gr
 import outlines
@@ -13,44 +14,24 @@ import requests
 from outlines import models, generate, samplers
 from pydantic import BaseModel
 
-
-def merge_games(clues, num_merges=10):
-    """Generates around 10 merges of words from the given clues.
-
-    Args:
-      clues: A list of clues, where each clue is a list containing the words, the answer, and the explanation.
-      num_merges: The approximate number of merges to generate (default: 10).
-
-    Returns:
-      A list of tuples, where each tuple contains the merged words and the indices of the selected rows.
-    """
-
-    merges = []
-    while len(merges) < num_merges:
-        num_rows = random.choice([3, 4])
-        selected_rows = random.sample(range(len(clues)), num_rows)
-        merged_words = " ".join([word for row in [clues[i][0] for i in selected_rows] for word in row])
-        if len(merged_words.split()) in [8, 9]:
-            merges.append((merged_words.split(), selected_rows))
-
-    return merges
-
+# Constants
+MAX_IMAGE_SIZE = (1024, 1024)
+TARGET_IMAGE_SIZE = 180_000
+NVIDIA_API_URL = "https://ai.api.nvidia.com/v1/vlm/microsoft/phi-3-vision-128k-instruct"
+MODEL_NAME = "microsoft/Phi-3-mini-4k-instruct"
 
 class Clue(BaseModel):
     word: str
     explanation: str
 
-
 class Group(BaseModel):
     words: List[str]
     clue: str
     explanation: str
 
-
 class Groups(BaseModel):
     groups: List[Group]
 
-
 example_clues = [
     (['ARROW', 'TIE', 'HONOR'], 'BOW', 'such as a bow and arrow, a bow tie, or a bow as a sign of honor'),
     (['DOG', 'TREE'], 'BARK', 'such as the sound a dog makes, or a tree is made of bark'),
@@ -81,8 +62,36 @@ example_clues = [
      'such as the Arctic being home to seals, or shutting a seal on an envelope, or a stamp being a type of seal'),
 ]
 
+def create_random_word_groups(clues: List[Tuple[List[str], str, str]], target_groups: int = 10) -> List[Tuple[List[str], List[int]]]:
+    """
+    Creates approximately 'target_groups' random groups of words from the given clues.
 
-def group_words(words):
+    Args:
+        clues: A list of clues, where each clue is a tuple (words, answer, explanation).
+        target_groups: The desired number of groups to create.
+
+    Returns:
+        A list of tuples, each containing a list of merged words and their corresponding indices.
+    """
+    groups = []
+    while len(groups) < target_groups:
+        num_rows = random.choice([3, 4])
+        selected_indices = random.sample(range(len(clues)), num_rows)
+        merged_words = [word for row in [clues[i][0] for i in selected_indices] for word in row]
+        if len(merged_words) in [8, 9]:
+            groups.append((merged_words, selected_indices))
+    return groups
+
+def group_words(word_list: List[str]) -> List[Group]:
+    """
+    Groups the given words into 3 to 4 thematic groups.
+
+    Args:
+        word_list: A list of words to be grouped.
+
+    Returns:
+        A list of Group objects representing the grouped words.
+    """
     @outlines.prompt
     def chat_group_template(system_prompt, query, history=[]):
         '''<s><|system|>
@@ -99,10 +108,11 @@ def group_words(words):
         '''
 
     grouping_system_prompt = ("You are an assistant for the game Codenames. Your task is to help players by grouping a "
-                              "given group of secrets into 3 to 4 groups. Each group should consist of secrets that "
-                              "share a common theme or other word connections such as homonym, hypernyms or synonyms")
+                              "given set of words into 3 to 4 groups. Each group should consist of words that "
+                              "share a common theme or other word connections such as homonyms, hypernyms, or synonyms.")
+    
     example_groupings = []
-    merges = merge_games(example_clues, 5)
+    merges = create_random_word_groups(example_clues, 5)
     for merged_words, indices in merges:
         groups = [{
             "secrets": example_clues[i][0],
@@ -111,20 +121,25 @@ def group_words(words):
         } for i in indices]
         example_groupings.append((merged_words, json.dumps(groups, separators=(',', ':'))))
 
-    prompt = chat_group_template(grouping_system_prompt, words, example_groupings)
+    prompt = chat_group_template(grouping_system_prompt, word_list, example_groupings)
     sampler = samplers.greedy()
     generator = generate.json(model, Groups, sampler)
 
-    print("Grouping words:", words)
-    generations = generator(
-        prompt,
-        max_tokens=500
-    )
-    print("Got groupings: ", generations)
+    print(f"Grouping words: {word_list}")
+    generations = generator(prompt, max_tokens=500)
+    print(f"Generated groupings: {generations}")
     return generations.groups
 
+def generate_clue(group: List[str]) -> Clue:
+    """
+    Generates a single-word clue for the given group of words.
+
+    Args:
+        group: A list of words to generate a clue for.
 
-def generate_clues(group):
+    Returns:
+        A Clue object containing the generated word and its explanation.
+    """
     @outlines.prompt
     def chat_clue_template(system, query, history=[]):
         '''<s><|system|>
@@ -140,23 +155,29 @@ def generate_clues(group):
         <|assistant|>
         '''
 
-    clue_system_prompt = ("You are a codenames game companion. Your task is to give a single word clue related to "
-                          "a given group of words. You will only respond with a single word clue. Compound words are "
+    clue_system_prompt = ("You are a Codenames game companion. Your task is to give a single word clue related to "
+                          "a given group of words. Respond with a single word clue only. Compound words are "
                           "allowed. Do not include the word 'Clue'. Do not provide explanations or notes.")
 
     prompt = chat_clue_template(clue_system_prompt, group, example_clues)
-    # sampler = samplers.greedy()
     sampler = samplers.multinomial(2, top_k=10)
     generator = generate.json(model, Clue, sampler)
     generations = generator(prompt, max_tokens=100)
-    print("Got clues: ", generations)
+    print(f"Generated clues: {generations}")
     return generations[0]
 
-
-def jpeg_with_target_size(im, target):
-    """Return the image as JPEG with the given name at best quality that makes less than "target" bytes
+def compress_image_to_jpeg(image: 'PIL.Image', target_size: int) -> bytes:
+    """
+    Compresses the image to JPEG format with the best quality that fits within the target size.
 
     https://stackoverflow.com/a/52281257
+
+    Args:
+        image: The PIL Image object to compress.
+        target_size: The target file size in bytes.
+
+    Returns:
+        The compressed image as bytes.
     """
     # Min and Max quality
     qmin, qmax = 25, 96
@@ -167,65 +188,90 @@ def jpeg_with_target_size(im, target):
 
         # Encode into memory and get size
         buffer = io.BytesIO()
-        im.save(buffer, format="JPEG", quality=m)
+        image.save(buffer, format="JPEG", quality=m)
         s = buffer.getbuffer().nbytes
 
-        if s <= target:
+        if s <= target_size:
             qacc = m
             qmin = m + 1
-        elif s > target:
+        elif s > target_size:
             qmax = m - 1
 
     # Write to disk at the defined quality
     if qacc > -1:
         image_byte_array = io.BytesIO()
-        print("Acceptable quality", im, im.format, f"{im.size}x{im.mode}")
-        im.save(image_byte_array, format='JPEG', quality=qacc)
+        print("Acceptable quality", image, image.format, f"{image.size}x{image.mode}")
+        image.save(image_byte_array, format='JPEG', quality=qacc)
         return image_byte_array.getvalue()
 
+def process_image(img: 'PIL.Image') -> gr.update:
+    """
+    Processes the uploaded image to detect words for the Codenames game.
 
-def process_image(img):
-    # Resize the image
-    max_size = (1024, 1024)
-    img.thumbnail(max_size)
+    Args:
+        img: The uploaded PIL Image object.
 
-    image_byte_array = jpeg_with_target_size(img, 180_000)
+    Returns:
+        A gradio update object with the detected words.
+    """
+    img.thumbnail(MAX_IMAGE_SIZE)
+    image_byte_array = compress_image_to_jpeg(img, TARGET_IMAGE_SIZE)
     image_b64 = base64.b64encode(image_byte_array).decode()
 
-    invoke_url = "https://ai.api.nvidia.com/v1/vlm/microsoft/phi-3-vision-128k-instruct"
-    stream = False
-
-    if os.environ.get("NVIDIA_API_KEY", "").startswith("nvapi-"):
-        print("Valid NVIDIA_API_KEY already in the environment. Delete to reset")
-
     headers = {
         "Authorization": f"Bearer {os.environ.get('NVIDIA_API_KEY', '')}",
-        "Accept": "text/event-stream" if stream else "application/json"
+        "Accept": "application/json"
     }
 
     payload = {
         "messages": [
             {
                 "role": "user",
-                "content": f'Identify the words in this game of Codenames. Provide only a list of words. Provide the '
-                           f'words in capital letters only. <img src="data:image/png;base64,{image_b64}" />'
+                "content": f'Identify the words in this game of Codenames. Provide only a list of words in capital letters. <img src="data:image/png;base64,{image_b64}" />'
             }
         ],
         "max_tokens": 512,
         "temperature": 0.1,
         "top_p": 0.70,
-        "stream": stream
+        "stream": False
     }
 
-    response = requests.post(invoke_url, headers=headers, json=payload)
+    response = requests.post(NVIDIA_API_URL, headers=headers, json=payload)
     if response.ok:
         print(response.json())
-        # Define the pattern to match uppercase words separated by commas
         pattern = r'[A-Z]+(?:\s+[A-Z]+)?'
         words = re.findall(pattern, response.json()['choices'][0]['message']['content'])
-
         return gr.update(choices=words, value=words)
 
+def pad_or_truncate_groups(groups: List[Group], target_length: int = 4) -> List[Group]:
+    """
+    Ensures the list of groups has exactly target_length elements, padding with empty Groups if necessary.
+
+    Args:
+        groups: The list of Group objects to pad or truncate.
+        target_length: The desired length of the list.
+
+    Returns:
+        A list of Group objects with the specified length.
+    """
+    truncated_groups = groups[:target_length]
+    return truncated_groups + [Group(words=[], clue='', explanation='') for _ in range(target_length - len(truncated_groups))]
+
+def group_words_callback(words: List[str]) -> List[gr.update]:
+    """
+    Callback function to group the selected words.
+
+    Args:
+        words: A list of words to group.
+
+    Returns:
+        A list of gradio update objects for each group input.
+    """
+    groups = group_words(words)
+    groups = pad_or_truncate_groups(groups, 4)
+    print(f"Generated groups: {groups}")
+    return [gr.update(value=group.words, choices=group.words, info=group.explanation) for group in groups]
+
 
 if __name__ == '__main__':
     with gr.Blocks() as demo:
@@ -235,10 +281,7 @@ if __name__ == '__main__':
         with gr.Row():
             game_image = gr.Image(type="pil")
             word_list_input = gr.Dropdown(label="Enter list of words (comma separated)",
-                                          choices='WEREWOLF, CHAIN, MOSQUITO, CRAFT, RANCH, LIP, VALENTINE, CLOUD, '
-                                                  'BEARD, BUNK, SECOND, SADDLE, BUCKET, JAIL, ANT, POCKET, LACE, '
-                                                  'BREAK, CUCKOO, FLAT, NIL, TIN, CHERRY, CHRISTMAS, MOSES, '
-                                                  'TEAM'.split(', '),
+                                          choices=[],
                                           multiselect=True,
                                           interactive=True)
 
@@ -277,7 +320,7 @@ if __name__ == '__main__':
 
         def generate_clues_callback(group):
             print("Generating clues: ", group)
-            g = generate_clues(group)
+            g = generate_clue(group)
             return gr.update(value=g.word, info=g.explanation)