Update app.py

app.py

@@ -1,14 +1,10 @@
 import os
-import spaces
-
-from threading import Thread
-from typing import Iterator, List, Tuple
-import json
 import requests
-
 import gradio as gr
 import torch
-import transformers
+import spaces
+from threading import Thread
+from typing import Iterator, List, Tuple
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
 # Description for the Gradio Interface
@@ -36,17 +32,20 @@ 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
+    try:
+        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
+    except Exception as e:
+        print(f"Error loading model {model_id}: {e}")
 
 # Function to log comparisons
 def log_comparison(model1_name: str, model2_name: str, question: str, answer1: str, answer2: str, winner: str = None):
@@ -57,7 +56,6 @@ def log_comparison(model1_name: str, model2_name: str, question: str, answer1: s
         "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:
@@ -70,14 +68,17 @@ def log_comparison(model1_name: str, model2_name: str, question: str, answer1: s
 # 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,
-    )
+    try:
+        inputs = tokenizer(
+            [x[1] for x in chat_history] + [message],
+            return_tensors="pt",
+            truncation=True,
+            padding=True,
+            max_length=MAX_INPUT_TOKEN_LENGTH,
+        )
+    except Exception as e:
+        print(f"Error preparing input for model {model_id}: {e}")
+        inputs = tokenizer([message], return_tensors="pt", padding=True, max_length=MAX_INPUT_TOKEN_LENGTH)
     return inputs
 
 # Function to generate responses from models
@@ -101,24 +102,28 @@ def generate(
         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)
+    try:
+        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)
+    except Exception as e:
+        print(f"Error generating response from model {model_id}: {e}")
+        yield "Error generating response."
 
 # Function to compare two models
 def compare_models(
@@ -135,15 +140,20 @@ def compare_models(
         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)))
+    try:
+        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))
+        chat_history1.append((message, output1))
+        chat_history2.append((message, output2))
 
-    log_comparison(model1_name, model2_name, message, output1, output2)
+        log_comparison(model1_name, model2_name, message, output1, output2)
 
-    return chat_history1, chat_history2, chat_history1, chat_history2
+        return chat_history1, chat_history2, chat_history1, chat_history2
+    except Exception as e:
+        print(f"Error comparing models: {e}")
+        error_message = [("System", "Error comparing models.")]
+        return error_message, error_message, chat_history1, chat_history2
 
 # Function to log the voting result
 def vote_better(model1_name, model2_name, question, answer1, answer2, choice):