Update app.py

app.py

@@ -206,7 +206,10 @@ def train_model():
     return "\n\n".join(metric_str_all), [os.path.join("./VRM-Emotions", file) for file in os.listdir("./VRM-Emotions")]
 
 async def train_model_async(progress=gr.Progress(track_tqdm=True)):
-    return await asyncio.to_thread(train_model)
+    try:
+        return await asyncio.to_thread(train_model)
+    except Exception as e:
+        raise gr.Error(e)
 
 DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、常に日本語で回答してください。"
 memory = [{"role": "system", "content": DEFAULT_SYSTEM_PROMPT}]
@@ -217,69 +220,75 @@ def reset_memory():
     return None
 
 def init():
-    global llamaTokenizer, llamaModel
-    llamaTokenizer = AutoTokenizer.from_pretrained("elyza/Llama-3-ELYZA-JP-8B", trust_remote_code=True)
-    llamaModel = AutoModelForCausalLM.from_pretrained(
-        "elyza/Llama-3-ELYZA-JP-8B",
-        torch_dtype="auto",
-        device_map="auto",
-    )
-    llamaModel.eval()
+    try:
+        global llamaTokenizer, llamaModel
+        llamaTokenizer = AutoTokenizer.from_pretrained("elyza/Llama-3-ELYZA-JP-8B", trust_remote_code=True)
+        llamaModel = AutoModelForCausalLM.from_pretrained(
+            "elyza/Llama-3-ELYZA-JP-8B",
+            torch_dtype="auto",
+            device_map="auto",
+        )
+        llamaModel.eval()
 
-    return [gr.Button(visible=False), gr.Textbox(visible=True), gr.Button(visible=True)]
+        return [gr.Button(visible=False), gr.Textbox(visible=True), gr.Button(visible=True)]    
+    except Exception as e:
+        raise gr.Error(e)
 
 async def chat(message):
-    async with Translator() as translator:
-        translated_input = await translator.translate(message, dest="ja")
-        jp_input = translated_input.text
-        output_language = translated_input.src
-
-        global memory
+    try:
+        async with Translator() as translator:
+            translated_input = await translator.translate(message, dest="ja")
+            jp_input = translated_input.text
+            output_language = translated_input.src
 
-        memory.append({"role": "user", "content": jp_input})
+            global memory
 
-        prompt = llamaTokenizer.apply_chat_template(
-            memory,
-            tokenize=False,
-            add_generation_prompt=True
-        )
+            memory.append({"role": "user", "content": jp_input})
 
-        token_ids = llamaTokenizer.encode(
-            prompt, add_special_tokens=False, return_tensors="pt"
-        )
+            prompt = llamaTokenizer.apply_chat_template(
+                memory,
+                tokenize=False,
+                add_generation_prompt=True
+            )
 
-        with torch.no_grad():
-            output_ids = llamaModel.generate(
-                token_ids.to(llamaModel.device),
-                max_new_tokens=1024,
-                do_sample=True,
-                temperature=0.6,
-                top_p=0.9,
+            token_ids = llamaTokenizer.encode(
+                prompt, add_special_tokens=False, return_tensors="pt"
             )
 
-        output = llamaTokenizer.decode(
-            output_ids.tolist()[0][token_ids.size(1):], skip_special_tokens=True
-        )
+            with torch.no_grad():
+                output_ids = llamaModel.generate(
+                    token_ids.to(llamaModel.device),
+                    max_new_tokens=1024,
+                    do_sample=True,
+                    temperature=0.6,
+                    top_p=0.9,
+                )
 
-        memory.append({"role": "assistant", "content": output})
+            output = llamaTokenizer.decode(
+                output_ids.tolist()[0][token_ids.size(1):], skip_special_tokens=True
+            )
 
-        translated_output = await translator.translate(output, dest=output_language)
+            memory.append({"role": "assistant", "content": output})
 
-        if os.path.exists(f"./VRM-Emotions/model.safetensors"):
-            newTokenizer = AutoTokenizer.from_pretrained("./VRM-Emotions", trust_remote_code=True)
-            newConfig = AutoConfig.from_pretrained("./VRM-Emotions")
-            newModel = AutoModelForSequenceClassification.from_pretrained("./VRM-Emotions", config=newConfig)
+            translated_output = await translator.translate(output, dest=output_language)
 
-            sentence = newTokenizer(output, return_tensors="pt", padding=True, truncation=True, max_length=512)
-            with torch.no_grad():
-                outputs = newModel(**sentence)
-                logits = outputs.logits
-                predicted = torch.argmax(logits, dim=-1).item()
-                label = newModel.config.id2label[predicted]
+            if os.path.exists(f"./VRM-Emotions/model.safetensors"):
+                newTokenizer = AutoTokenizer.from_pretrained("./VRM-Emotions", trust_remote_code=True)
+                newConfig = AutoConfig.from_pretrained("./VRM-Emotions")
+                newModel = AutoModelForSequenceClassification.from_pretrained("./VRM-Emotions", config=newConfig)
 
-            return "[" + label + "] " + translated_output.text
-        else:
-            return translated_output.text
+                sentence = newTokenizer(output, return_tensors="pt", padding=True, truncation=True, max_length=512)
+                with torch.no_grad():
+                    outputs = newModel(**sentence)
+                    logits = outputs.logits
+                    predicted = torch.argmax(logits, dim=-1).item()
+                    label = newModel.config.id2label[predicted]
+
+                return "[" + label + "] " + translated_output.text
+            else:
+                return translated_output.text
+    except Exception as e:
+        raise gr.Error(e)
 
 with gr.Blocks() as demo:
     with gr.Tab("Prepare Dataset"):