Update app.py

app.py

@@ -10,10 +10,10 @@ MODEL_MAPPING = {
     "bart-large-cnn": ("facebook/bart-large-cnn", AutoModelForSeq2SeqLM),
     "t5-base": ("t5-base", AutoModelForSeq2SeqLM),
     "distilbert-base-uncased": ("distilbert-base-uncased", AutoModelForCausalLM),
-    "bert-base-uncased": ("bert-base-uncased", AutoModelForSeq2SeqLM),
-    "roberta-base": ("roberta-base", AutoModelForSeq2SeqLM),
-    "xlnet-base-cased": ("xlnet-base-cased", AutoModelForSeq2SeqLM),
-    "albert-base-v2": ("albert-base-v2", AutoModelForSeq2SeqLM),
+    "bert-base-uncased": ("bert-base-uncased", AutoModelForCausalLM),
+    "roberta-base": ("roberta-base", AutoModelForCausalLM),
+    "xlnet-base-cased": ("xlnet-base-cased", AutoModelForCausalLM),
+    "albert-base-v2": ("albert-base-v2", AutoModelForCausalLM),
 }
 
 def get_model_and_tokenizer(model_name):
@@ -41,4 +41,13 @@ model_name = gr.inputs.Dropdown(choices=list(MODEL_MAPPING.keys()), label="Model
 response = gr.outputs.Textbox(label="Generated Response")
 
 # Define Gradio interface
-gr.Interface(fn=generate_response, inputs=[input_text, model_name], outputs=response).launch()
+gr.Interface(fn=generate_response, inputs=[input_text, model_name], outputs=response).launch()
+
+# Define error handler
+def error_handler(e):
+    print(e)
+    with open("errors.txt", "a") as f:
+        f.write(str(e) + "\n")
+
+# Set error handler
+gr.set_error_handler(error_handler)