Update app.py

app.py

@@ -5,27 +5,31 @@ import json
 import requests
 
 import gradio as gr
-import spaces
 import torch
 import transformers
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
+# Description for the Gradio Interface
 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
+# Constants
+MAX_MAX_NEW_TOKENS = 256
+DEFAULT_MAX_NEW_TOKENS = 128
 MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
 
+# Device configuration
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
+# Model options
 MODEL_OPTIONS = [
     "sarvamai/OpenHathi-7B-Hi-v0.1-Base",
     "TokenBender/Navarna_v0_1_OpenHermes_Hindi"
 ]
 
+# Load models and tokenizers
 models = {}
 tokenizers = {}
 
@@ -42,6 +46,7 @@ for model_id in MODEL_OPTIONS:
     if tokenizers[model_id].pad_token_id is None:
         tokenizers[model_id].pad_token_id = tokenizers[model_id].eos_token_id
 
+# Function to log comparisons
 def log_comparison(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None):
     log_data = {
         "question": question,
@@ -60,13 +65,27 @@ def log_comparison(model1_name: str, model2_name: str, question: str, answer1: s
     except requests.RequestException as e:
         print(f"Error sending log to server: {e}")
 
-@spaces.GPU(duration=90)
+# Function to prepare input
+def prepare_input(model_id: str, message: str, chat_history: List[Tuple[str, str]]):
+    tokenizer = tokenizers[model_id]
+    # Prepare inputs for the model
+    inputs = tokenizer(
+        [x[1] for x in chat_history] + [message],
+        return_tensors="pt",
+        truncation=True,
+        padding=True,
+        max_length=MAX_INPUT_TOKEN_LENGTH,
+    )
+    return inputs
+
+# Function to generate responses from models
+@spaces.GPU(duration=120)
 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,
+    temperature: float = 0.4,
     top_p: float = 0.95,
 ) -> Iterator[str]:
     model = models[model_id]
@@ -99,6 +118,7 @@ def generate(
         outputs.append(text)
         yield "".join(outputs)
 
+# Function to compare two models
 def compare_models(
     model1_name: str,
     model2_name: str,
@@ -123,13 +143,13 @@ def compare_models(
 
     return chat_history1, chat_history2, chat_history1, chat_history2
 
-
-
+# Function to log the voting result
 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."
 
+# Gradio UI setup
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(DESCRIPTION)
     
@@ -174,9 +194,6 @@ with gr.Blocks(css="style.css") as demo:
         outputs=[vote_output]
     )
 
+# Main function to run the Gradio app
 if __name__ == "__main__":
-    # Start Flask server in a separate thread
-
-    
-    # Start Gradio app with public link
-    demo.queue(max_size=3).launch(share=True)
+    demo.queue(max_size=3).launch(share=True)