Update app.py

app.py

@@ -223,9 +223,9 @@ def build_models(model_type, config, enable_optimization=False):
     # build stdit
     # we load model from HuggingFace directly so that we don't need to
     # handle model download logic in HuggingFace Space
-    from transformers import AutoModelForSeq2SeqLM
+    from transformers import AutoModel
 
-    stdit = AutoModelForSeq2SeqLM.from_pretrained(
+    stdit = AutoModel.from_pretrained(
         HF_STDIT_MAP[model_type],
         enable_flash_attn=enable_optimization,
         trust_remote_code=True,