Update app.py

app.py

@@ -1,5 +1,9 @@
 import gradio as gr
 from transformers import pipeline
+from transformers import AutoTokenizer, TraXL
+
+tokenizer = AutoTokenizer.from_pretrained("Wonder-Griffin/TraXL")
+ # model = TraXL.from_pretrained("Wonder-Griffin/TraXL")
 
 pipe = pipeline("text-generation", model="Wonder-Griffin/TraXL")
 
@@ -7,7 +11,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 = AutoModelForCasualLM.from_pretrained(
+model = TraXL.from_pretrained(
     model_path,
     device_map="auto",
     torch_dtype='auto'