Add device support

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+import torch
 from ui import title, description
 from langs import LANGS
 
@@ -9,6 +10,7 @@ CKPT = "facebook/nllb-200-distilled-600M"
 model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
 tokenizer = AutoTokenizer.from_pretrained(CKPT)
 
+device = 0 if torch.cuda.is_available() else -1
 
 def translate(text, src_lang, tgt_lang, max_length=400):
     """
@@ -19,7 +21,8 @@ def translate(text, src_lang, tgt_lang, max_length=400):
                                     tokenizer=tokenizer,
                                     src_lang=src_lang,
                                     tgt_lang=tgt_lang,
-                                    max_length=max_length)
+                                    max_length=max_length,
+                                    device=device)
 
     result = translation_pipeline(text)
     return result[0]['translation_text']

requirements.txt

@@ -1,3 +1,3 @@
-transformers
+git+https://github.com/huggingface/transformers
 gradio
 torch