Update app.py

app.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
     pipeline,
-    AutoModel,
+    AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteriaList
@@ -82,10 +82,9 @@ class S3ModelLoader:
                 s3_uri, local_files_only=True
             )
             
-            model = AutoModel.from_pretrained(
+            model = AutoModelForCausalLM.from_pretrained(
                 s3_uri, config=config, local_files_only=True
             )
-            
                 
             tokenizer = AutoTokenizer.from_pretrained(
                 s3_uri, config=config, local_files_only=True
@@ -106,7 +105,7 @@ class S3ModelLoader:
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
                 
-                model = AutoModel.from_pretrained(
+                model = AutoModelForCausalLM.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )