working version

app.py

@@ -0,0 +1,299 @@
+import base64
+import io
+import math
+import os
+import random
+import json
+import re
+from typing import List
+
+import gradio as gr
+import outlines
+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
+
+
+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'),
+    (['MONEY', 'RIVER', 'ROB', 'BLOOD'], 'CRIME', 'such as money being stolen, a river being a potential crime scene, '
+                                                  'robbery, or blood being a result of a violent crime'),
+    (['BEEF', 'TURKEY', 'FIELD', 'GRASS'], 'GROUND',
+     'such as ground beef, a turkey being a ground-dwelling bird, a field or grass being a type of ground'),
+    (['BANK', 'GUITAR', 'LIBRARY'], 'NOTE',
+     'such as a bank note, a musical note on a guitar, or a note being a written comment in a library book'),
+    (['ROOM', 'PIANO', 'TYPEWRITER'], 'KEYS', 'such as a room key, piano keys, or typewriter keys'),
+    (['TRAFFIC', 'RADAR', 'PHONE'], 'SIGNAL', 'such as traffic signals, radar signals, or phone signals'),
+    (['FENCE', 'PICTURE', 'COOKIE'], 'FRAME',
+     'such as a frame around a yard, a picture frame, or a cookie cutter being a type of frame'),
+    (['YARN', 'VIOLIN', 'DRESS'], 'STRING', 'strings like material, instrument, clothing fastener'),
+    (['JUMP', 'FLOWER', 'CLOCK'], 'SPRING',
+     'such as jumping, flowers blooming in the spring, or a clock having a sprint component'),
+    (['SPY', 'KNIFE'], 'WAR',
+     'Both relate to aspects of war, such as spies being involved in war or knives being used as weapons'),
+    (['STADIUM', 'SHOE', 'FIELD'], 'SPORT', 'Sports like venues, equipment, playing surfaces'),
+    (['TEACHER', 'CLUB'], 'SCHOOL',
+     'such as a teacher being a school staff member or a club being a type of school organization'),
+    (['CYCLE', 'ARMY', 'COURT', 'FEES'], 'CHARGE', 'charges like electricity, battle, legal, payments'),
+    (['FRUIT', 'MUSIC', 'TRAFFIC', 'STUCK'], 'JAM',
+     'Jams such as fruit jam, a music jam session, traffic jam, or being stuck in a jam'),
+    (['POLICE', 'DOG', 'THIEF'], 'CRIME',
+     'such as police investigating crimes, dogs being used to detect crimes, or a thief committing a crime'),
+    (['ARCTIC', 'SHUT', 'STAMP'], 'SEAL',
+     'such as the Arctic being home to seals, or shutting a seal on an envelope, or a stamp being a type of seal'),
+]
+
+
+def group_words(words):
+    @outlines.prompt
+    def chat_group_template(system_prompt, query, history=[]):
+        '''<s><|system|>
+        {{ system_prompt }}
+        {% for example in history %}
+        <|user|>
+        {{ example[0] }}<|end|>
+        <|assistant|>
+        {{ example[1] }}<|end|>
+        {% endfor %}
+        <|user|>
+        {{ query }}<|end|>
+        <|assistant|>
+        '''
+
+    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")
+    example_groupings = []
+    merges = merge_games(example_clues, 5)
+    for merged_words, indices in merges:
+        groups = [{
+            "secrets": example_clues[i][0],
+            "clue": example_clues[i][1],
+            "explanation": example_clues[i][2]
+        } for i in indices]
+        example_groupings.append((merged_words, json.dumps(groups, separators=(',', ':'))))
+
+    prompt = chat_group_template(grouping_system_prompt, words, 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)
+    return generations.groups
+
+
+def generate_clues(group):
+    @outlines.prompt
+    def chat_clue_template(system, query, history=[]):
+        '''<s><|system|>
+        {{ system }}
+        {% for example in history %}
+        <|user|>
+        {{ example[0] }}<|end|>
+        <|assistant|>
+        {"Clue": "{{ example[1] }}", "Description": "{{ example[2] }}" }<|end|>
+        {% endfor %}
+        <|user|>
+        {{ query }}<|end|>
+        <|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 "
+                          "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)
+    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
+
+    https://stackoverflow.com/a/52281257
+    """
+    # Min and Max quality
+    qmin, qmax = 25, 96
+    # Highest acceptable quality found
+    qacc = -1
+    while qmin <= qmax:
+        m = math.floor((qmin + qmax) / 2)
+
+        # Encode into memory and get size
+        buffer = io.BytesIO()
+        im.save(buffer, format="JPEG", quality=m)
+        s = buffer.getbuffer().nbytes
+
+        if s <= target:
+            qacc = m
+            qmin = m + 1
+        elif s > target:
+            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)
+        return image_byte_array.getvalue()
+
+
+def process_image(img):
+    # Resize the image
+    max_size = (1024, 1024)
+    img.thumbnail(max_size)
+
+    image_byte_array = jpeg_with_target_size(img, 180_000)
+    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"
+    }
+
+    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}" />'
+            }
+        ],
+        "max_tokens": 512,
+        "temperature": 0.1,
+        "top_p": 0.70,
+        "stream": stream
+    }
+
+    response = requests.post(invoke_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)
+
+
+if __name__ == '__main__':
+    with gr.Blocks() as demo:
+        gr.Markdown("# *Codenames* clue generator")
+        gr.Markdown("Provide a list of words to generate a clue")
+
+        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(', '),
+                                          multiselect=True,
+                                          interactive=True)
+
+        with gr.Row():
+            detect_words_button = gr.Button("Detect Words")
+            group_words_button = gr.Button("Group Words")
+
+        dropdowns, buttons, outputs = [], [], []
+
+        for i in range(4):
+            with gr.Row():
+                group_input = gr.Dropdown(label=f"Group {i + 1}",
+                                        choices=[],
+                                        allow_custom_value=True,
+                                        multiselect=True,
+                                        interactive=True)
+                clue_button = gr.Button("Generate Clue", size='sm')
+                clue_output = gr.Textbox(label=f"Clue {i + 1}")
+                dropdowns.append(group_input)
+                buttons.append(clue_button)
+                outputs.append(clue_output)
+
+        def pad_or_truncate(lst, n=4):
+            # Ensure the length of the list is at most n
+            truncated_lst = lst[:n]
+            return truncated_lst + (n - len(truncated_lst)) * [Group(words=[],clue='',explanation='')]
+
+
+        def group_words_callback(words):
+            groups = group_words(words)
+            groups = pad_or_truncate(groups, 4)
+            print("Got groups: ", groups, type(groups))
+            return [gr.update(value=groups[i].words, choices=groups[i].words, info=groups[i].explanation) for i in range(4)]
+
+
+
+        def generate_clues_callback(group):
+            print("Generating clues: ", group)
+            g = generate_clues(group)
+            return gr.update(value=g.word, info=g.explanation)
+
+
+        model = models.transformers("microsoft/Phi-3-mini-4k-instruct",
+                                    model_kwargs={'device_map': "cuda", 'torch_dtype': "auto",
+                                                  'trust_remote_code': True,
+                                                  'attn_implementation': "flash_attention_2"})
+
+        detect_words_button.click(fn=process_image,
+                                  inputs=game_image,
+                                  outputs=[word_list_input])
+        group_words_button.click(fn=group_words_callback,
+                                 inputs=word_list_input,
+                                 outputs=dropdowns)
+
+        for i in range(4):
+            buttons[i].click(generate_clues_callback, inputs=dropdowns[i], outputs=outputs[i])
+
+    demo.launch(share=False)