Update app.py

update model class to AutoModelForCausalLM

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import os
 import torch
 import spaces
-from transformers import AutoModel
+from transformers import AutoModelForCausalLM
 from huggingface_hub import login
 
 hf_token = os.getenv("HF_TOKEN")
@@ -14,7 +14,7 @@ def get_model_summary(model_name):
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     
     # Load the model and move it to the selected device
-    model = AutoModel.from_pretrained(model_name, trust_remote_code=True).to(device)
+    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True).to(device)
     
     # Return the model's architecture as a string
     return str(model)