Update app.py

app.py

@@ -10,22 +10,23 @@ MODELS = ["facebook/nllb-200-distilled-600M","DigitalUmuganda/Finetuned-NLLB"]
 # tokenizer = AutoTokenizer.from_pretrained(CKPT)
 
 device = 0 if torch.cuda.is_available() else -1
+fb_model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-600M")
+du_model = AutoModelForSeq2SeqLM.from_pretrained("DigitalUmuganda/Finetuned-NLLB")
+tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-600M")
 
-
+models = {"facebook/nllb-200-distilled-600M":fb_model,"DigitalUmuganda/Finetuned-NLLB":du_model}
 # def translate(text, src_lang, tgt_lang, max_length=400):
 def translate(text, src_lang, tgt_lang, CKPT, max_length=400):
 
     """
     Translate the text from source lang to target lang
     """
-    model = AutoModelForSeq2SeqLM.from_pretrained(CKPT)
-    tokenizer = AutoTokenizer.from_pretrained(CKPT)
-
+  
     translation_pipeline = pipeline(TASK,
                                     tokenizer=tokenizer,
                                     src_lang=src_lang,
                                     tgt_lang=tgt_lang,
-                                    model = model,
+                                    model = models[CKPT],
                                     max_length=max_length,
                                     device=device)