import gradio as gr from response_db import ResponseDb from create_cache import Game_Cache import numpy as np from PIL import Image import pandas as pd import torch import pickle import uuid import nltk nltk.download('punkt') db = ResponseDb() css = """ .chatbot {display:flex;flex-direction:column} .msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%} .msg.user {background-color:cornflowerblue;color:white;align-self:self-end} .msg.bot {background-color:lightgray} .na_button {background-color:red;color:red} """ from model.run_question_asking_model import return_modules, return_modules_yn question_model, response_model_simul, _, caption_model = return_modules() question_model_yn, response_model_simul_yn, _, caption_model_yn = return_modules_yn() class Game_Session: def __init__(self, taskid, yn, hard_setting): self.yn = yn self.hard_setting = hard_setting global question_model, response_model_simul, caption_model global question_model_yn, response_model_simul_yn, caption_model_yn self.question_model = question_model self.response_model_simul = response_model_simul self.caption_model = caption_model self.question_model_yn = question_model_yn self.response_model_simul_yn = response_model_simul_yn self.caption_model_yn = caption_model_yn global image_files, images_np, p_y_x, p_r_qy, p_y_xqr, captions, questions, target_questions self.image_files, self.image_np, self.p_y_x, self.p_r_qy, self.p_y_xqr = None, None, None, None, None self.captions, self.questions, self.target_questions = None, None, None self.history = [] self.game_id = str(uuid.uuid4()) self.set_curr_models() def set_curr_models(self): if self.yn: self.curr_question_model, self.curr_caption_model, self.curr_response_model_simul = self.question_model_yn, self.caption_model_yn, self.response_model_simul_yn else: self.curr_question_model, self.curr_caption_model, self.curr_response_model_simul = self.question_model, self.caption_model, self.response_model_simul def get_next_question(self): return self.curr_question_model.select_best_question(self.p_y_x, self.questions, self.images_np, self.captions, self.curr_response_model_simul) def ask_a_question(input, taskid, gs): gs.history.append(input) gs.p_r_qy = gs.curr_response_model_simul.get_p_r_qy(input, gs.history[-2], gs.images_np, gs.captions) gs.p_y_xqr = gs.p_y_x*gs.p_r_qy gs.p_y_xqr = gs.p_y_xqr/torch.sum(gs.p_y_xqr)if torch.sum(gs.p_y_xqr) != 0 else torch.zeros_like(gs.p_y_xqr) gs.p_y_x = gs.p_y_xqr gs.questions.remove(gs.history[-2]) db.add(gs.game_id, taskid, len(gs.history)//2-1, gs.history[-2], gs.history[-1]) gs.history.append(gs.get_next_question()) top_prob = torch.max(gs.p_y_x).item() top_pred = torch.argmax(gs.p_y_x).item() if top_prob > 0.8: gs.history = gs.history[:-1] db.add(gs.game_id, taskid, len(gs.history)//2, f"Guess: Image {top_pred}", "") # write some HTML html = "
The model identified Image {top_pred+1} as the image. Please select a new task ID to continue.
" return html, gs, gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False) else: if not gs.yn: return html, gs, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False) else: return html, gs, gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=True), gr.Button.update(visible=True) def set_images(taskid): pilot_study = pd.read_csv("pilot-study.csv") taskid_original = taskid taskid = pilot_study['mscoco-id'].tolist()[int(taskid)] with open(f'cache/{int(taskid)}.p', 'rb') as fp: game_cache = pickle.load(fp) gs = Game_Session(int(taskid), game_cache.yn, game_cache.hard_setting) id1 = f"./mscoco-images/val2014/{game_cache.image_files[0]}" id2 = f"./mscoco-images/val2014/{game_cache.image_files[1]}" id3 = f"./mscoco-images/val2014/{game_cache.image_files[2]}" id4 = f"./mscoco-images/val2014/{game_cache.image_files[3]}" id5 = f"./mscoco-images/val2014/{game_cache.image_files[4]}" id6 = f"./mscoco-images/val2014/{game_cache.image_files[5]}" id7 = f"./mscoco-images/val2014/{game_cache.image_files[6]}" id8 = f"./mscoco-images/val2014/{game_cache.image_files[7]}" id9 = f"./mscoco-images/val2014/{game_cache.image_files[8]}" id10 = f"./mscoco-images/val2014/{game_cache.image_files[9]}" gs.image_files = [id1, id2, id3, id4, id5, id6, id7, id8, id9, id10] gs.image_files = [x[15:] for x in gs.image_files] gs.images_np = [np.asarray(Image.open(f"./mscoco-images/{i}")) for i in gs.image_files] gs.images_np = [np.dstack([i]*3) if len(i.shape)==2 else i for i in gs.images_np] gs.p_y_x = (torch.ones(10)/10).to(gs.curr_question_model.device) gs.captions = gs.curr_caption_model.get_captions(gs.image_files) gs.questions, gs.target_questions = gs.curr_question_model.get_questions(gs.image_files, gs.captions, 0) gs.curr_question_model.reset_question_bank() gs.curr_question_model.question_bank = game_cache.question_dict first_question = gs.curr_question_model.select_best_question(gs.p_y_x, gs.questions, gs.images_np, gs.captions, gs.curr_response_model_simul) first_question_html = f"Current Task ID: {int(taskid_original)}
" if not gs.yn: return id1, id2, id3, id4, id5, id6, id7, id8, id9, id10, gs, first_question_html, gr.HTML.update(value=html, visible=True), gr.Textbox.update(visible=True, value=''), gr.Button.update(visible=True), gr.Button.update(visible=True), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False) else: return id1, id2, id3, id4, id5, id6, id7, id8, id9, id10, gs, first_question_html, gr.HTML.update(value=html, visible=True), gr.Textbox.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=False), gr.Number.update(visible=False), gr.Button.update(visible=False), gr.Button.update(visible=True), gr.Button.update(visible=True) with gr.Blocks(title="Image Q&A Guessing Game", css=css) as demo: gr.HTML("\
Imagine you are playing 20-questions with an AI model.
\
The AI model plays the role of the question asker. You play the role of the responder.
\
There are 10 images. Your image is Image 1. The other images are distraction images.\
The model can see all 10 images and all the questions and answers for the current set of images. It will ask a question based on the available information.
\
The goal of the model is to accurately guess the correct image (i.e. Image 1) in as few turns as possible.
\
Your goal is to help the model guess the image by answering as clearly and accurately as possible.
\
Guidelines:
\