Update app.py

app.py

@@ -2,15 +2,18 @@ import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import pytesseract
 
-# Initialize chat model (You can change the model here)
-chat_model = pipeline("text-generation", model="gpt2")  # You can switch to any model of your choice
 
 
+# Load model directly
+tokenizer = AutoTokenizer.from_pretrained("aubmindlab/bert-base-arabic")
+model = AutoModelForCausalLM.from_pretrained("aubmindlab/bert-base-arabic")
+
 # Chat function
 def chat_fn(history, user_input):
     conversation = {"history": history, "user": user_input}
-    response = chat_model(user_input, max_length=50, num_return_sequences=1)
-    conversation["bot"] = response[0]['generated_text']
+    # Use the model for Arabic
+    response = model.generate(input_ids=tokenizer.encode(user_input, return_tensors="pt"), max_length=50)
+    conversation["bot"] = tokenizer.decode(response[0], skip_special_tokens=True)
     history.append((user_input, conversation["bot"]))
     return history, ""