pre-loading models

app.py

@@ -10,7 +10,7 @@ language_map = {'English':'en','Swahili':'sw','Fon':'fon','Igbo':'ig',
                 'Efik':'efi','Twi':'twi','Afrikaans':'af','Yoruba':'yo'}  
                 
 #List of available languages I worked on.
-available_languages =list(language_map.keys())
+available_language_pairs =['en-sw','en-af','en-ar','efi-en','en-ha','en-ig','en-fon','en-twi','sn-en','sw-en','yo-en']
 
 
 def load_config(path="configs/default.yaml") -> dict:
@@ -27,8 +27,8 @@ def load_config(path="configs/default.yaml") -> dict:
     return cfg
   
 def load_model(source_language_,target_language_):  
-    source_language = language_map[source_language_]
-    target_language = language_map[target_language_]
+    #source_language = language_map[source_language_]
+    #target_language = language_map[target_language_]
     
     #source_language = 'en'
     #target_language = 'sw'
@@ -50,11 +50,23 @@ def load_model(source_language_,target_language_):
     params = load_params_for_prediction(parsed_yaml_file,best_ckpt)
     return params
 
+#Load models of all available language pairs
+model_mapping = {} 
+for availabe_lang in available_language_pairs:
+    model_mapping.update({availabe_lang:load_model(availabe_lang.split('-')[0],availabe_lang.split('-')[1])})
+
+
+
 def get_translation(source_language,target_language,source_sentence=None,source_file=None):
     '''
     This takes a sentence and gets the translation.
     type_=2 tells joeynmt translate that it should expect a sentence. 
     '''
+    
+    source_language_ = language_map[source_language]
+    target_language_ = language_map[target_language]
+    params = 
+    
     source = source_sentence
     type_=2
     if source_file!=None:
@@ -62,13 +74,14 @@ def get_translation(source_language,target_language,source_sentence=None,source_
         source = source_file.name
     try:
        
-        params = load_model(source_language,target_language)
+        #params = load_model(source_language,target_language)
+        params = model_mapping[f'{source_language_}-{target_language_}']
         pred = translate(params,source,type_)
     except Exception:
         return f'There was an issue loading the translation model for {source_language} -> {target_language}. Try another pair please'
     
     return pred[0] if source_file==None else pred
-    
+