Spaces:
Runtime error
Runtime error
| import random | |
| import json | |
| import os | |
| import httpx | |
| from gradio_client import Client | |
| from urllib.parse import urljoin | |
| import gspread | |
| from oauth2client.service_account import ServiceAccountCredentials | |
| import gradio as gr | |
| from huggingface_hub import login | |
| import asyncio | |
| from concurrent.futures import ThreadPoolExecutor | |
| BOOK_MAPPING = { | |
| "Spaceman on a Spree": "In the context of \"Spaceman on a Spree\", written by Mack Reynolds in 1961,", | |
| "Charity Case": "In the context of \"Charity Case\", written by Jim Harmon in 1972,", | |
| "A Gift from Earth": "In the context of \"A Gift From Earth\", written by Manly Banister in 1950,", | |
| "Pick a Crime": "In the context of \"Pick a Crime\", written by Richard Rein Smith in 1970,", | |
| "Dangerous Quarry": "In the context of \"Dangerous Quarry\", written by Jim Harmon in 1972,", | |
| "Lost in Translation": "In the context of \"Lost in Translation\", written by Larry M. Harris in 1972," | |
| } | |
| CATEGORY_MAPPING = { | |
| "Identity": "identity", | |
| "Motivation": "motivation", | |
| "Relationship": "relationship", | |
| "Event": "event", | |
| } | |
| MODEL_VARIANTS = ["rephrase", "rephrase_summarize", "entigraph"] | |
| model_responses = {"Model_A": "", "Model_B": "", "Model_C": ""} | |
| class ModelManager: | |
| def __init__(self, books): | |
| self.book_model_assignments = {} | |
| self.assign_models_to_books(books) | |
| self.model_A = Client("mep296/llama-3-8b-rephrase-quality") | |
| self.model_B = Client("mep296/llama-3-8b-rephrase-summarize-quality") | |
| self.model_C = Client("mep296/llama-3-8b-entigraph-quality") | |
| self.template_text = self._load_template() | |
| def _load_template(self): | |
| with open("prompt_template.txt", "r", encoding="utf-8") as file: | |
| return file.read() | |
| async def get_model_response_async(self, model_name, client, prompt): | |
| try: | |
| formatted_prompt = self.template_text.format(prompt) | |
| loop = asyncio.get_running_loop() | |
| with ThreadPoolExecutor() as executor: | |
| response = await loop.run_in_executor( | |
| executor, | |
| client.predict, | |
| formatted_prompt, | |
| "/predict" | |
| ) | |
| model_responses[model_name] = response | |
| return response | |
| except (httpx.ReadTimeout, httpx.ConnectTimeout) as e: | |
| print(f"Timeout while getting response from {model_name}: {str(e)}") | |
| return f"⚠️ Model {model_name} is waking up... This may take a few minutes. Please wait and try again shortly. ⏳" | |
| except Exception as e: | |
| print(f"Error getting response from {model_name}: {str(e)}") | |
| return f"Error: Could not get response from {model_name}. Please try again in a few minutes." | |
| async def get_all_model_responses_async(self, prompt): | |
| tasks = [ | |
| self.get_model_response_async("Model_A", self.model_A, prompt), | |
| self.get_model_response_async("Model_B", self.model_B, prompt), | |
| self.get_model_response_async("Model_C", self.model_C, prompt) | |
| ] | |
| responses = await asyncio.gather(*tasks) | |
| return responses | |
| def get_book_model_mapping(self, book): | |
| if book not in self.book_model_assignments: | |
| shuffled_models = random.sample(MODEL_VARIANTS, len(MODEL_VARIANTS)) | |
| self.book_model_assignments[book] = { | |
| "Model A": shuffled_models[0], | |
| "Model B": shuffled_models[1], | |
| "Model C": shuffled_models[2] | |
| } | |
| return self.book_model_assignments[book] | |
| def assign_models_to_books(self, books): | |
| for book in books: | |
| self.get_book_model_mapping(book) | |
| class SheetManager: | |
| def __init__(self): | |
| scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"] | |
| creds = ServiceAccountCredentials.from_json_keyfile_dict(variables_keys, scope) | |
| client = gspread.authorize(creds) | |
| self.sheet = client.open("Model Blind Comparison Ratings").sheet1 | |
| def append_rating(self, rating_data): | |
| self.sheet.append_row([ | |
| rating_data["book"], | |
| rating_data["category"], | |
| rating_data["prompt"], | |
| rating_data["rephrase_rating"], | |
| rating_data["rephrase_summarize_rating"], | |
| rating_data["entigraph_rating"], | |
| rating_data["rephrase_response"], | |
| rating_data["rephrase_summarize_response"], | |
| rating_data["entigraph_response"], | |
| ]) | |
| class ModelComparisonApp: | |
| def __init__(self): | |
| self.model_manager = ModelManager(BOOK_MAPPING.keys()) | |
| self.sheet_manager = SheetManager() | |
| self.votes = [] | |
| self.selected_book = "Spaceman on a Spree" | |
| self.selected_book_string = BOOK_MAPPING["Spaceman on a Spree"] | |
| self.selected_category_string = "" | |
| self.chat_history_A = [] | |
| self.chat_history_B = [] | |
| self.chat_history_C = [] | |
| self.state = gr.State(value="") | |
| self.loop = asyncio.get_event_loop() | |
| def create_interface(self): | |
| text_size = gr.themes.sizes.text_lg | |
| with gr.Blocks(theme=gr.themes.Default(text_size=text_size), fill_width=True) as demo: | |
| gr.Markdown("# Model Blind Comparison") | |
| with gr.Group(): | |
| with gr.Row(): | |
| chat_interfaces = self._create_chat_interfaces() | |
| with gr.Row(): | |
| ratings = self._create_rating_sliders() | |
| with gr.Row(equal_height=True): | |
| submit_button = gr.Button(value="⭐ Submit Ratings", interactive=False) | |
| submission_status = gr.Textbox(label="Submission Status", interactive=False) | |
| with gr.Row(equal_height=True): | |
| input_elements = self._create_input_elements() | |
| self._setup_event_handlers(demo, chat_interfaces, ratings, submit_button, submission_status, input_elements) | |
| return demo | |
| def _create_chat_interfaces(self): | |
| interfaces = {} | |
| for model in ['A', 'B', 'C']: | |
| interfaces[model] = gr.Chatbot( | |
| getattr(self, f'chat_history_{model}'), | |
| type="messages", | |
| label=f"Model {model}", | |
| height=650, | |
| show_copy_button=True | |
| ) | |
| return interfaces | |
| def _create_rating_sliders(self): | |
| return { | |
| str(i): gr.Slider(1, 5, step=1, label=f"Rate Response {chr(64+i)}", | |
| interactive=True, value=3) | |
| for i in range(1, 4) | |
| } | |
| def _create_input_elements(self): | |
| return { | |
| 'book': gr.Dropdown(choices=list(BOOK_MAPPING.keys()), | |
| label="Select a Book", interactive=True, scale=1), | |
| 'category': gr.Dropdown(choices=list(CATEGORY_MAPPING.keys()), | |
| label="Select a Question Category", interactive=True, scale=1), | |
| 'question': gr.Textbox(label="Type a Question", max_lines=1, | |
| placeholder="e.g. What is the relationship between Harry Potter and Sirius Black?", | |
| interactive=True, scale=2), | |
| 'send': gr.Button("Send", scale=0, variant="primary", interactive=False) | |
| } | |
| def respond(self, message): | |
| if not message.strip(): | |
| raise gr.Error("Message cannot be empty!") | |
| prompt = f"{self.selected_book_string} {message}" | |
| mapping_dict = self.model_manager.book_model_assignments[self.selected_book] | |
| model_order = ["rephrase", "rephrase_summarize", "entigraph"] | |
| model_to_index = {model: i for i, model in enumerate(model_order)} | |
| responses = self.loop.run_until_complete( | |
| self.model_manager.get_all_model_responses_async(prompt) | |
| ) | |
| chats = [] | |
| for response in responses: | |
| chat = [] | |
| chat.append({"role": "user", "content": prompt}) | |
| chat.append({"role": "assistant", "content": response}) | |
| chats.append(chat) | |
| reordered_chats = [chats[model_to_index[mapping_dict[model]]] for model in ["Model A", "Model B", "Model C"]] | |
| return reordered_chats | |
| def get_votes(self, book, category, question, rating_A, rating_B, rating_C): | |
| model_mapping = self.model_manager.get_book_model_mapping(book) | |
| rating_data = { | |
| "book": book, | |
| "category": category, | |
| "prompt": question, | |
| "rephrase_rating": rating_A if model_mapping["Model A"] == "rephrase" else | |
| rating_B if model_mapping["Model B"] == "rephrase" else rating_C, | |
| "rephrase_summarize_rating": rating_A if model_mapping["Model A"] == "rephrase_summarize" else | |
| rating_B if model_mapping["Model B"] == "rephrase_summarize" else rating_C, | |
| "entigraph_rating": rating_A if model_mapping["Model A"] == "entigraph" else | |
| rating_B if model_mapping["Model B"] == "entigraph" else rating_C, | |
| "rephrase_response": model_responses["Model_A"], | |
| "rephrase_summarize_response": model_responses["Model_B"], | |
| "entigraph_response": model_responses["Model_C"] | |
| } | |
| self.votes.append(rating_data) | |
| self.sheet_manager.append_rating(rating_data) | |
| return ("Ratings submitted successfully!", gr.update(interactive=False)) | |
| def _setup_event_handlers(self, demo, chat_interfaces, ratings, submit_button, submission_status, input_elements): | |
| def enable_send_btn(book, category, question): | |
| return gr.update(interactive=bool(book and category and question)) | |
| def enable_button_group(model_A, model_B, model_C): | |
| return gr.update(interactive=bool(model_A and model_B and model_C)) | |
| def update_selected_book(book_selection): | |
| self.selected_book = book_selection | |
| self.selected_book_string = BOOK_MAPPING.get(book_selection, "") | |
| return self.selected_book_string | |
| for input_name in ['book', 'category', 'question']: | |
| input_elements[input_name].change( | |
| enable_send_btn, | |
| inputs=[input_elements['book'], input_elements['category'], input_elements['question']], | |
| outputs=[input_elements['send']] | |
| ) | |
| input_elements['book'].change( | |
| update_selected_book, | |
| inputs=[input_elements['book']], | |
| outputs=[self.state] | |
| ) | |
| submit_button.click( | |
| self.get_votes, | |
| inputs=[input_elements['book'], input_elements['category'], input_elements['question'], | |
| ratings['1'], ratings['2'], ratings['3']], | |
| outputs=[submission_status, submit_button] | |
| ) | |
| input_elements['send'].click( | |
| self.respond, | |
| inputs=[input_elements['question']], | |
| outputs=list(chat_interfaces.values()) | |
| ) | |
| for interface in chat_interfaces.values(): | |
| interface.change( | |
| enable_button_group, | |
| inputs=list(chat_interfaces.values()), | |
| outputs=[submit_button] | |
| ) | |
| if __name__ == "__main__": | |
| PRIVATE_KEY = os.environ.get('PRIVATE_KEY').replace('\\n', '\n') | |
| PRIVATE_KEY_ID = os.environ.get('PRIVATE_KEY_ID').replace('\\n', '\n') | |
| variables_keys = { | |
| "type": "service_account", | |
| "project_id": "summer-presence-450117-r7", | |
| "private_key_id": PRIVATE_KEY_ID, | |
| "private_key": PRIVATE_KEY, | |
| "client_email": "model-blind-comparison@summer-presence-450117-r7.iam.gserviceaccount.com", | |
| "client_id": "117681363507032419648", | |
| "auth_uri": "https://accounts.google.com/o/oauth2/auth", | |
| "token_uri": "https://oauth2.googleapis.com/token", | |
| "auth_provider_x509_cert_url": "https://www.googleapis.com/oauth2/v1/certs", | |
| "client_x509_cert_url": "https://www.googleapis.com/robot/v1/metadata/x509/model-blind-comparison%40summer-presence-450117-r7.iam.gserviceaccount.com", | |
| "universe_domain": "googleapis.com" | |
| } | |
| app = ModelComparisonApp() | |
| demo = app.create_interface() | |
| demo.launch() |