import spaces import base64 import io import math import os import random import json import re from typing import List, Tuple import PIL import gradio as gr import outlines import requests from outlines import models, generate, samplers from pydantic import BaseModel # Install Flash attention import subprocess subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True) # 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'), (['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 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. 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 @spaces.GPU(duration=120) 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=[]): '''<|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 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 = create_random_word_groups(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, word_list, example_groupings) sampler = samplers.greedy() generator = generate.json(model, Groups, sampler) print(f"Grouping words: {word_list}") generations = generator(prompt, max_tokens=500) print(f"Generated groupings: {generations}") return generations.groups @spaces.GPU(duration=120) 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. Returns: A Clue object containing the generated word and its explanation. """ @outlines.prompt def chat_clue_template(system, query, history=[]): '''<|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. 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.multinomial(2, top_k=10) generator = generate.json(model, Clue, sampler) generations = generator(prompt, max_tokens=100) print(f"Generated clues: {generations}") return generations[0] 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 # Highest acceptable quality found qacc = -1 while qmin <= qmax: m = math.floor((qmin + qmax) / 2) # Encode into memory and get size buffer = io.BytesIO() image.save(buffer, format="JPEG", quality=m) s = buffer.getbuffer().nbytes if s <= target_size: qacc = m qmin = m + 1 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", 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. Args: img: The uploaded PIL Image object. 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() headers = { "Authorization": f"Bearer {os.environ.get('NVIDIA_API_KEY', '')}", "Accept": "application/json" } payload = { "messages": [ { "role": "user", "content": f'Identify the words in this game of Codenames. Provide only a list of words in capital letters. ' } ], "max_tokens": 512, "temperature": 0.1, "top_p": 0.70, "stream": False } response = requests.post(NVIDIA_API_URL, headers=headers, json=payload) if response.ok: print(response.json()) 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] def generate_clues_callback(group): """ Callback function to generate a clue for the given words. Args: words: a list of words. Returns: A gradio update object for the given group input. """ print("Generating clues: ", group) g = generate_clue(group) return gr.update(value=g.word, info=g.explanation) 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=[], 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) 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)