import os from threading import Thread from typing import Iterator, List, Tuple import json import requests import gradio as gr import spaces import torch import transformers from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer DESCRIPTION = """\ # Zero GPU Model Comparison Arena Select two different models from the dropdowns and see how they perform on the same input. """ MAX_MAX_NEW_TOKENS = 1024 DEFAULT_MAX_NEW_TOKENS = 256 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096")) device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") MODEL_OPTIONS = [ "sarvamai/OpenHathi-7B-Hi-v0.1-Base", "TokenBender/Navarna_v0_1_OpenHermes_Hindi" ] models = {} tokenizers = {} for model_id in MODEL_OPTIONS: tokenizers[model_id] = AutoTokenizer.from_pretrained(model_id) models[model_id] = AutoModelForCausalLM.from_pretrained( model_id, device_map="auto", load_in_8bit=True, ) models[model_id].eval() # Set pad_token_id to eos_token_id if it's not set if tokenizers[model_id].pad_token_id is None: tokenizers[model_id].pad_token_id = tokenizers[model_id].eos_token_id def log_comparison(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None): log_data = { "question": question, "model1": {"name": model1_name, "answer": answer1}, "model2": {"name": model2_name, "answer": answer2}, "winner": winner } # Send log data to remote server try: response = requests.post('http://144.24.151.32:5000/log', json=log_data, timeout=5) if response.status_code == 200: print("Successfully logged to server") else: print(f"Failed to log to server. Status code: {response.status_code}") except requests.RequestException as e: print(f"Error sending log to server: {e}") @spaces.GPU(duration=90) def generate( model_id: str, message: str, chat_history: List[Tuple[str, str]], max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, temperature: float = 0.7, top_p: float = 0.95, ) -> Iterator[str]: model = models[model_id] tokenizer = tokenizers[model_id] inputs = prepare_input(model_id, message, chat_history) input_ids = inputs.input_ids if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH: input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:] gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.") input_ids = input_ids.to(model.device) streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True) generate_kwargs = dict( input_ids=input_ids, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=True, top_p=top_p, temperature=temperature, num_beams=1, pad_token_id=tokenizer.eos_token_id, ) t = Thread(target=model.generate, kwargs=generate_kwargs) t.start() outputs = [] for text in streamer: outputs.append(text) yield "".join(outputs) def compare_models( model1_name: str, model2_name: str, message: str, chat_history1: List[Tuple[str, str]], chat_history2: List[Tuple[str, str]], max_new_tokens: int, temperature: float, top_p: float, ) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]], List[Tuple[str, str]], List[Tuple[str, str]]]: if model1_name == model2_name: error_message = [("System", "Error: Please select two different models.")] return error_message, error_message, chat_history1, chat_history2 output1 = "".join(list(generate(model1_name, message, chat_history1, max_new_tokens, temperature, top_p))) output2 = "".join(list(generate(model2_name, message, chat_history2, max_new_tokens, temperature, top_p))) chat_history1.append((message, output1)) chat_history2.append((message, output2)) log_comparison(model1_name, model2_name, message, output1, output2) return chat_history1, chat_history2, chat_history1, chat_history2 def vote_better(model1_name, model2_name, question, answer1, answer2, choice): winner = model1_name if choice == "Model 1" else model2_name log_comparison(model1_name, model2_name, question, answer1, answer2, winner) return f"You voted that {winner} performs better. This has been logged." with gr.Blocks(css="style.css") as demo: gr.Markdown(DESCRIPTION) with gr.Row(): with gr.Column(): model1_dropdown = gr.Dropdown(choices=MODEL_OPTIONS, label="Model 1", value=MODEL_OPTIONS[0]) chatbot1 = gr.Chatbot(label="Model 1 Output") with gr.Column(): model2_dropdown = gr.Dropdown(choices=MODEL_OPTIONS, label="Model 2", value=MODEL_OPTIONS[1]) chatbot2 = gr.Chatbot(label="Model 2 Output") text_input = gr.Textbox(label="Input Text", lines=3) with gr.Row(): max_new_tokens = gr.Slider(label="Max new tokens", minimum=1, maximum=MAX_MAX_NEW_TOKENS, value=DEFAULT_MAX_NEW_TOKENS) temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, value=0.7) top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, value=0.95) compare_btn = gr.Button("Compare Models") with gr.Row(): better1_btn = gr.Button("Model 1 is Better") better2_btn = gr.Button("Model 2 is Better") vote_output = gr.Textbox(label="Voting Result") compare_btn.click( compare_models, inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, max_new_tokens, temperature, top_p], outputs=[chatbot1, chatbot2, chatbot1, chatbot2] ) better1_btn.click( vote_better, inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, gr.Textbox(value="Model 1", visible=False)], outputs=[vote_output] ) better2_btn.click( vote_better, inputs=[model1_dropdown, model2_dropdown, text_input, chatbot1, chatbot2, gr.Textbox(value="Model 2", visible=False)], outputs=[vote_output] ) if __name__ == "__main__": # Start Flask server in a separate thread # Start Gradio app with public link demo.queue(max_size=3).launch(share=True)