Update app.py

app.py

@@ -1,6 +1,7 @@
 from transformers import pipeline
 from datasets import Dataset
 import streamlit as st
+import torch
 
 # Set the background color and layout with set_page_config
 st.set_page_config(
@@ -21,12 +22,12 @@ text_input = st.text_area("Enter English text to translate", height=150, value=s
 # Define your model from Hugging Face
 model_directory = "repleeka/eng-tagin-nmt"
 
-# Initialize the pipeline for translation
+device = 0 if torch.cuda.is_available() else -1
 translation_pipeline = pipeline(
-    task="translation_en_to_tagin",  # Adjust task if needed
-    model=model_directory,
-    tokenizer=model_directory,
-    device=0  # Use GPU (if available)
+    task="translation",
+    model="repleeka/eng-tagin-nmt",
+    tokenizer="repleeka/eng-tagin-nmt",
+    device=device
 )
 
 # Translate button