Update app.py

app.py

@@ -5,6 +5,7 @@ import os
 import httpx
 
 logging.set_verbosity_error()
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 def download_argos_model(from_code, to_code):
     import argostranslate.package
@@ -53,7 +54,8 @@ if model_name == 'Helsinki-NLP':
         model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 if model_name.startswith('t5'):
     tokenizer = T5Tokenizer.from_pretrained(model_name)
-    model = T5ForConditionalGeneration.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
+
 
 st.write("Selected language combination:", sselected_language, " - ", tselected_language, "Selected model:", model_name)
 submit_button = st.button("Translate")
@@ -77,7 +79,7 @@ if submit_button:
     elif model_name.startswith('t5'):
         prompt = f'translate {sselected_language} to {tselected_language}: {input_text}'
         print(prompt)
-        input_ids = tokenizer.encode(prompt, return_tensors='pt')
+        input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
         # Perform translation
         output_ids = model.generate(input_ids)
         # Decode the translated text
@@ -104,8 +106,6 @@ if submit_button:
             translated_text = f"No Argos model for {sselected_language} to {tselected_language}. Try other model or languages combination!"
         except Exception as error:
             translated_text = error
-            # download_argos_model(sl, tl)
-            # translated_text = argostranslate.translate.translate(input_text, sl, tl)
     # Display the translated text
     print(translated_text)
     st.write(f"Translated text from {sselected_language} to {tselected_language} using {model_name}:")