Update backend_utils.py

backend_utils.py

@@ -8,6 +8,12 @@ import pickle
 import torch
 from sklearn.multiclass import OneVsRestClassifier
 from sklearn.ensemble import RandomForestClassifier
+import spacy
+
+# nlp = spacy.load("en_core_web_trf")
+nlp = spacy.load("en_core_web_sm")
+
+
 
 class wrappedTokenizer(RobertaTokenizer):
     def __call__(self, text_input):
@@ -24,7 +30,7 @@ def generate_index(db):
         })
     return index_list
 
-def load_db(db_metadata_path, db_constructor_path, db_params_path):
+def load_db(db_metadata_path, db_constructor_path, db_params_path, exclusion_list_path):
     '''
     Function to load dataframe
 
@@ -42,7 +48,11 @@ def load_db(db_metadata_path, db_constructor_path, db_params_path):
     db_constructor.dropna(inplace=True)
     db_params = pd.read_csv(db_params_path)
     db_params.dropna(inplace=True)
-    return db_metadata, db_constructor, db_params
+    with open(exclusion_list_path, 'r') as f:
+        ex_list = f.read()
+    ex_list = ex_list.split("\n")
+
+    return db_metadata, db_constructor, db_params, ex_list
 
 
 
@@ -177,8 +187,9 @@ def retrieve_libraries(retrieval_model, model_input, db_metadata):
     '''
     results = retrieval_model(model_input)
     library_ids = [item.get('id') for item in results]
+    scores = [item.get('similarity') for item in results]
     library_names = [id_to_libname(item, db_metadata) for item in library_ids]
-    return library_ids, library_names
+    return library_ids, library_names, scores
 
 def prepare_input_generative_model(library_ids, db_constructor):
     '''
@@ -423,10 +434,11 @@ def initialize_all_components(config):
     classifier_head: a random forest model
     '''
     # load db
-    db_metadata, db_constructor, db_params = load_db(
+    db_metadata, db_constructor, db_params, ex_list = load_db(
         config.get('db_metadata_path'), 
         config.get('db_constructor_path'),
-        config.get('db_params_path')
+        config.get('db_params_path'),
+        config.get('exclusion_list_path')
     )
 
     # load model
@@ -443,14 +455,14 @@ def initialize_all_components(config):
         config.get('classifier_head_path')
     )
 
-    return db_metadata, db_constructor, db_params, model_retrieval, model_generative, tokenizer_generative, model_classifier, classifier_head, tokenizer_classifier
+    return db_metadata, db_constructor, db_params, ex_list, model_retrieval, model_generative, tokenizer_generative, model_classifier, classifier_head, tokenizer_classifier
 
 def make_predictions(input_query, 
     model_retrieval, 
     model_generative,  
     model_classifier, classifier_head,
     tokenizer_generative, tokenizer_classifier,
-    db_metadata, db_constructor, db_params,
+    db_metadata, db_constructor, db_params, ex_list,
     config):
     '''
     Function to retrieve relevant libraries, generate API usage patterns, and predict the hw configs
@@ -467,9 +479,28 @@ def make_predictions(input_query,
     Returns:
     predictions (list): a list of dictionary containing the prediction details
     '''
-    library_ids, library_names = retrieve_libraries(model_retrieval, input_query, db_metadata)
+    print("retrieve libraries")
+    queries = extract_keywords(input_query.lower(), ex_list)
+
+    temp_list = []
+    for query in queries:
+        temp_library_ids, temp_library_names, temp_scores = retrieve_libraries(model_retrieval, query, db_metadata)
+
+        if len(temp_library_ids) > 0:
+            for id_, name, score in zip(temp_library_ids, temp_library_names, temp_scores):
+                temp_list.append((id_, name, score))
+    
+    library_ids = []
+    library_names = []
+    if len(temp_list) > 0:
+        sorted_list = sorted(temp_list, key=lambda tup: tup[2], reverse=True)
+        sorted_list = sorted_list[:config.get('max_k')]
+        for item in sorted_list:
+            library_ids.append(item[0])
+            library_names.append(item[1])
 
     if len(library_ids) == 0:
+        print("null libraries")
         return "null"
 
     print("generate usage patterns")
@@ -500,4 +531,49 @@ def make_predictions(input_query,
     print("finished the predictions")
     predictions = get_metadata_library(predictions, db_metadata)
 
-    return predictions
+    return predictions
+
+def extract_series(x):
+    name = x.replace("-", " ").replace("_", " ")
+    name = name.split()
+    series = []
+    for token in name:
+        if token.isalnum() and not(token.isalpha()) and not(token.isdigit()):
+            series.append(token)
+    if len(series) > 0:
+        return series
+    else:
+        return [x]
+
+def extract_keywords(query, ex_list):
+    doc = nlp(query)
+    keyword_candidates = []
+
+    # extract keywords
+    for chunk in doc.noun_chunks:
+        temp_list = []
+
+        for token in chunk:
+            if token.text not in ex_list and token.pos_ not in ("DET", "PRON", "CCONJ", "NUM"):
+                temp_list.append(token.text)
+
+        if len(temp_list) > 0:
+            keyword_candidates.append(" ".join(temp_list))
+
+    filtered_keyword_candidates = []
+    for keyword in keyword_candidates:
+        temp_candidates = extract_series(keyword)
+
+        for keyword in temp_candidates:
+
+            if len(keyword.split()) > 1:
+                doc = nlp(keyword)
+                for chunk in doc.noun_chunks:
+                    filtered_keyword_candidates.append(chunk.root.text)
+            else:
+                filtered_keyword_candidates.append(keyword)
+
+    if len(filtered_keyword_candidates) == 0:
+        filtered_keyword_candidates.append(query)
+    
+    return filtered_keyword_candidates