initial commit

app.py

@@ -0,0 +1,232 @@
+import json
+import os
+import random
+
+import google.generativeai as genai
+import gradio as gr
+from jinja2 import Environment, FileSystemLoader
+
+GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
+
+genai.configure(api_key=GOOGLE_API_KEY)
+
+model = genai.GenerativeModel("gemini-1.5-flash-latest",
+                              generation_config={"response_mime_type": "application/json"})
+
+
+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.
+    """
+
+    merged = []
+    while len(merged) < num_merges:
+        num_rows = random.choice([3, 4])
+        selected_rows = random.sample(range(len(clues)), num_rows)
+        words_list = [word for row in [clues[i][0] for i in selected_rows] for word in row]
+        if len(words_list) in [8, 9]:
+            merged.append((words_list, selected_rows))
+
+    return merged
+
+
+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'),
+]
+
+example_groupings = []
+merges = merge_games(example_clues, 5)
+for merged_words, indices in merges:
+    groups = [{
+        "words": 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=(',', ':'))))
+
+
+def render_jinja2_template(template, system_prompt, history, query):
+    env = Environment(loader=FileSystemLoader('.'))
+    template = env.from_string(template)
+    return template.render(system_prompt=system_prompt, history=history, query=query)
+
+
+def group_words(words):
+    template = '''
+    {% for example in history %}
+    INPUT:
+    {{ example[0] }}
+    OUTPUT:
+    {{ example[1] }}
+    {% endfor %}
+    INPUT:
+    {{ query }}
+    OUTPUT:
+
+    {{ system }}
+
+    Groups = {'words': list[str], 'clue': str, 'explanation': str}
+    Return: Groups
+    '''
+
+    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")
+
+    prompt = render_jinja2_template(template, grouping_system_prompt, example_groupings, words)
+    # print(prompt)
+    raw_response = model.generate_content(
+        prompt,
+        generation_config={'top_k': 3, 'temperature': 1.1})
+    response = json.loads(raw_response.text)
+
+    print("Grouping words:", words)
+
+    print("Got groupings: ", json.dumps(response, indent=4))
+    return [group["words"] for group in response]
+
+
+def generate_clues(group):
+    template = '''
+    {% for example in history %}
+    INPUT:
+    {{ example[0] }}
+    OUTPUT:
+    { 'clue':{{ example[1] }}, 'explanation':{{ example[2] }} }
+    {% endfor %}
+    INPUT:
+    {{ query }}
+    OUTPUT:
+
+
+    {{ system }}
+
+    Clue = {'clue': str, 'explanation': str}
+    Return: Clue
+    '''
+
+    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.")
+
+    prompt = render_jinja2_template(template, clue_system_prompt, example_clues, group)
+
+    raw_response = model.generate_content(
+        prompt,
+        generation_config={'top_k': 3, 'temperature': 1.1})
+    response = json.loads(raw_response.text)
+
+    print("Generating clues for: ", group)
+
+    print("Got clue: ", json.dumps(response, indent=4))
+    return response
+
+
+def process_image(img):
+    print(img, type(img))
+    raw_response = model.generate_content(['Identify the words in this game of Codenames. Provide only a list of '
+                                           'words. Provide the words in capital letters only. Group these words into '
+                                           '6 or 8 groups that can be guessed together using a single word clue in '
+                                           'the game of codenames. Give a response as json of the form: {"Game": '
+                                           '<list of words in the game>}', img], stream=True)
+    raw_response.resolve()
+    response = json.loads(raw_response.text)
+    words = response['Game']
+    return gr.update(choices=words, value=words)
+
+
+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)) * [None]
+
+
+    def group_words_callback(words):
+        word_groups = group_words(words)
+        word_groups = pad_or_truncate(word_groups, 4)
+        print("Got groups: ", word_groups, type(word_groups))
+        return [gr.update(value=word_groups[i], choices=words) for i in range(4)]
+
+
+    def generate_clues_callback(group):
+        print("Generating clues: ", group)
+        g = generate_clues(group)
+        return gr.update(value=g['clue'], info=g['explanation'])
+
+
+    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=True, debug=True)