0.12 catch exceptions

app.py

@@ -4,7 +4,9 @@ import torch
 import gradio as gr
 import logging
 from huggingface_hub import login
+
 import os
+import traceback
 
 from threading import Thread
 
@@ -12,6 +14,8 @@ from threading import Thread
 
 logging.basicConfig(level=logging.DEBUG)
 
+SPACER = '\n' + '*' * 40 + '\n'
+
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 login(token=HF_TOKEN)
 
@@ -51,97 +55,109 @@ def apply_chat_template(messages, add_generation_prompt=False):
 
 def load_model_a(model_id):
     global tokenizer_a, model_a, model_id_a
-    model_id_a = model_id # need to access model_id with tokenizer
-    tokenizer_a = AutoTokenizer.from_pretrained(model_id)
-    logging.debug(f"***** model A eos_token: {tokenizer_a.eos_token}")
-    model_a = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch_dtype,
-        device_map="auto",
-        trust_remote_code=True,
-    ).eval()
+    try: 
+        model_id_a = model_id # need to access model_id with tokenizer
+        tokenizer_a = AutoTokenizer.from_pretrained(model_id)
+        model_a = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch_dtype,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+    except Exception as e:
+        logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
+
     return gr.update(label=model_id)
     
+
 def load_model_b(model_id):
     global tokenizer_b, model_b, model_id_b
-    model_id_b = model_id
-    tokenizer_b = AutoTokenizer.from_pretrained(model_id)
-    logging.debug(f"***** model B eos_token: {tokenizer_b.eos_token}")
-    model_b = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch_dtype,
-        device_map="auto",
-        trust_remote_code=True,
-    ).eval()
-    model_b.tie_weights()
+    try: 
+        model_id_b = model_id
+        tokenizer_b = AutoTokenizer.from_pretrained(model_id)
+        logging.debug(f"***** model B eos_token: {tokenizer_b.eos_token}")
+        model_b = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch_dtype,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
+    except Exception as e:
+        logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
     return gr.update(label=model_id)
 
 @spaces.GPU()
 def generate_both(system_prompt, input_text, chatbot_a, chatbot_b, max_new_tokens=2048, temperature=0.2, top_p=0.9, repetition_penalty=1.1):
 
-    text_streamer_a = TextIteratorStreamer(tokenizer_a, skip_prompt=True)
-    text_streamer_b = TextIteratorStreamer(tokenizer_b, skip_prompt=True)
-
-    system_prompt_list = [{"role": "system", "content": system_prompt}] if system_prompt else []
-    input_text_list = [{"role": "user", "content": input_text}]
-
-    chat_history_a = []
-    for user, assistant in chatbot_a:
-        chat_history_a.append({"role": "user", "content": user})
-        chat_history_a.append({"role": "assistant", "content": assistant})
-
-    chat_history_b = []
-    for user, assistant in chatbot_b:
-        chat_history_b.append({"role": "user", "content": user})
-        chat_history_b.append({"role": "assistant", "content": assistant})
-    
-    new_messages_a = system_prompt_list + chat_history_a + input_text_list
-    new_messages_b = system_prompt_list + chat_history_b + input_text_list
-
-    input_ids_a = tokenizer_a.apply_chat_template(
-        new_messages_a,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(model_a.device)
-
-    input_ids_b = tokenizer_b.apply_chat_template(
-        new_messages_b,
-        add_generation_prompt=True,
-        return_tensors="pt"
-    ).to(model_b.device)
-
-    generation_kwargs_a = dict(
-        input_ids=input_ids_a,
-        streamer=text_streamer_a,
-        max_new_tokens=max_new_tokens,
-        pad_token_id=tokenizer_a.eos_token_id,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-    )
-    generation_kwargs_b = dict(
-        input_ids=input_ids_b,
-        streamer=text_streamer_b,
-        max_new_tokens=max_new_tokens,
-        pad_token_id=tokenizer_b.eos_token_id,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-    )
-
-    thread_a = Thread(target=model_a.generate, kwargs=generation_kwargs_a)
-    thread_b = Thread(target=model_b.generate, kwargs=generation_kwargs_b)
-
-    thread_a.start()
-    thread_b.start()
-
-    chatbot_a.append([input_text, ""])
-    chatbot_b.append([input_text, ""])
-
-    finished_a = False
-    finished_b = False
+    try: 
+        text_streamer_a = TextIteratorStreamer(tokenizer_a, skip_prompt=True)
+        text_streamer_b = TextIteratorStreamer(tokenizer_b, skip_prompt=True)
+
+        system_prompt_list = [{"role": "system", "content": system_prompt}] if system_prompt else []
+        input_text_list = [{"role": "user", "content": input_text}]
+
+        chat_history_a = []
+        for user, assistant in chatbot_a:
+            chat_history_a.append({"role": "user", "content": user})
+            chat_history_a.append({"role": "assistant", "content": assistant})
+
+        chat_history_b = []
+        for user, assistant in chatbot_b:
+            chat_history_b.append({"role": "user", "content": user})
+            chat_history_b.append({"role": "assistant", "content": assistant})
+        
+        new_messages_a = system_prompt_list + chat_history_a + input_text_list
+        new_messages_b = system_prompt_list + chat_history_b + input_text_list
+
+        input_ids_a = tokenizer_a.apply_chat_template(
+            new_messages_a,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to(model_a.device)
+
+        input_ids_b = tokenizer_b.apply_chat_template(
+            new_messages_b,
+            add_generation_prompt=True,
+            return_tensors="pt"
+        ).to(model_b.device)
+
+        logging.debug(f'model_a.device: {model_a.device}, model_b.device: {model_b.device}')
+
+        generation_kwargs_a = dict(
+            input_ids=input_ids_a,
+            streamer=text_streamer_a,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=tokenizer_a.eos_token_id,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+        
+        generation_kwargs_b = dict(
+            input_ids=input_ids_b,
+            streamer=text_streamer_b,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=tokenizer_b.eos_token_id,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+        )
+
+        thread_a = Thread(target=model_a.generate, kwargs=generation_kwargs_a)
+        thread_b = Thread(target=model_b.generate, kwargs=generation_kwargs_b)
+
+        thread_a.start()
+        thread_b.start()
+
+        chatbot_a.append([input_text, ""])
+        chatbot_b.append([input_text, ""])
+
+        finished_a = False
+        finished_b = False
+    except Exception as e:
+        logging.error(f'{SPACER} Error: {e}, Traceback {traceback.format_exc()}')
 
     while not (finished_a and finished_b):
         if not finished_a: