updatE

app.py

@@ -1,13 +1,17 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import os
+from transformers import pipeline, DiaForConditionalGeneration, AutoTokenizer
 
 # Load the model and tokenizer
 model_path = "EzioDevio/iraqi_dialect_llm"  # Use the Hugging Face Hub model path
 
+# Get token if available
+hf_token = os.getenv("HF_TOKEN")
+
 # Load the model and tokenizer directly from Hugging Face
-model = AutoModelForCausalLM.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = DiaForConditionalGeneration.from_pretrained(model_path, token=hf_token)
+tokenizer = AutoTokenizer.from_pretrained(model_path, token=hf_token)
+text_generator = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
 # Define the text generation function
 def generate_text(input_text):