added gpu support

app.py

@@ -1,6 +1,11 @@
 import gradio as gr
 from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 import os
+import torch
+import spaces
+
+device = "cuda:0" if torch.cuda.is_available() else "cpu"
+print(f'[INFO] Using device: {device}')
 
 # Define model paths
 MODEL_PATHS = {
@@ -14,7 +19,8 @@ MODEL_PATHS = {
 TOKEN = os.environ['TOKEN']
 
 # Translation function for Nano and Large models
-def translate_nano_large(text, model_path):
+@spaces.GPU
+def translate_nano_large(text, model_path, device='cuda:0'):
     translator = pipeline("translation", model=model_path, token=TOKEN)
     translated = translator(
         text, 
@@ -25,21 +31,24 @@ def translate_nano_large(text, model_path):
         do_sample=False,
         pad_token_id=translator.tokenizer.pad_token_id,
         bos_token_id=translator.tokenizer.bos_token_id,
-        eos_token_id=translator.tokenizer.eos_token_id
+        eos_token_id=translator.tokenizer.eos_token_id,
+        device=device,
     )
     return translated[0]["translation_text"]
 
 # Translation function for Ultra and Supreme models
-def translate_ultra_supreme(text, model_path):
+@spaces.GPU
+def translate_ultra_supreme(text, model_path, device='cuda:0'):
     model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token=TOKEN)
-    tokenizer = AutoTokenizer.from_pretrained(model_path, token=TOKEN)  # Remove src_lang and tgt_lang
+    tokenizer = AutoTokenizer.from_pretrained(model_path, src_lang="eng_Latn", tgt_lang="ary_Arab", token=TOKEN)
     translator = pipeline(
         "translation",
         model=model,
         tokenizer=tokenizer,
         max_length=512,
         src_lang="eng_Latn",  # Keep src_lang and tgt_lang in the pipeline
-        tgt_lang="ary_Arab"
+        tgt_lang="ary_Arab",
+        device=device,
     )
     translation = translator(text)[0]['translation_text']
     return translation