Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
 from transformers import pipeline
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModelForCasualLM
 
 tokenizer = AutoTokenizer.from_pretrained("Wonder-Griffin/TraXL")
-model = AutoModel.from_pretrained("Wonder-Griffin/TraXL")
 
 pipe = pipeline("text-generation", model="Wonder-Griffin/TraXL")
 
@@ -11,7 +10,7 @@ pipe = pipeline("text-generation", model="Wonder-Griffin/TraXL")
 model_path = "Wonder-Griffin/TraXL"  # Define this as needed
 
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModel.from_pretrained(
+model = AutoModelForCasualLM.from_pretrained(
     model_path,
     device_map="auto",
     torch_dtype='auto'