Update backend_utils.py

backend_utils.py

@@ -24,7 +24,7 @@ def generate_index(db):
         })
     return index_list
 
-def load_db(db_metadata_path, db_constructor_path):
+def load_db(db_metadata_path, db_constructor_path, db_params_path):
     '''
     Function to load dataframe
 
@@ -40,7 +40,9 @@ def load_db(db_metadata_path, db_constructor_path):
     db_metadata.dropna(inplace=True)
     db_constructor = pd.read_csv(db_constructor_path)
     db_constructor.dropna(inplace=True)
-    return db_metadata, db_constructor
+    db_params = pd.read_csv(db_params_path)
+    db_params.dropna(inplace=True)
+    return db_metadata, db_constructor, db_params
 
 
 
@@ -142,8 +144,6 @@ def get_metadata_library(predictions, db_metadata):
 
         else:
             prediction_dict['Description'] = "Description not found"
-        print(prediction_dict)
-        print("-----------------------------------------------------------------")
     return predictions_cp
 
 def id_to_libname(id_, db_metadata):
@@ -201,7 +201,7 @@ def prepare_input_generative_model(library_ids, db_constructor):
             )
     return output_dict
 
-def generate_api_usage_patterns(generative_model, tokenizer, model_input, num_beams, num_return_sequences):
+def generate_api_usage_patterns(generative_model, tokenizer, model_input, num_beams, num_return_sequences, max_length):
     '''
     Function to generate API usage patterns
 
@@ -221,7 +221,7 @@ def generate_api_usage_patterns(generative_model, tokenizer, model_input, num_be
         num_beams=num_beams,
         num_return_sequences=num_return_sequences, 
         early_stopping=True,
-        max_length=50
+        max_length=max_length
     )
     api_usage_patterns = tokenizer.batch_decode(
         model_output,
@@ -229,7 +229,36 @@ def generate_api_usage_patterns(generative_model, tokenizer, model_input, num_be
     )
     return api_usage_patterns
 
-def generate_api_usage_patterns_batch(generative_model, tokenizer, library_ids, db_constructor, num_beams, num_return_sequences):
+def add_params(api_usage_patterns, db_params, library_id):
+    patterns_cp = api_usage_patterns.copy()
+    valid = True
+    processed_sequences = []
+    for sequence in patterns_cp:
+        sequence_list = sequence.split()
+
+        if len(sequence_list) < 2:
+            continue
+
+        temp_list = []
+        for api in sequence_list:
+            temp_db = db_params[(db_params.id==library_id) & (db_params.methods==api.split(".")[-1])]
+
+            if len(temp_db) > 0:
+                param = temp_db.iloc[0].params
+                new_api = api + param
+                temp_list.append(new_api)
+            else:
+                valid = False
+                break
+        
+        if valid:
+            processed_sequences.append("[API-SEP]".join(temp_list))
+        else:
+            valid = True
+    return processed_sequences
+
+
+def generate_api_usage_patterns_batch(generative_model, tokenizer, library_ids, db_constructor, db_params, num_beams, num_return_sequences, max_length):
     '''
     Function to generate API usage patterns in batch
 
@@ -260,7 +289,8 @@ def generate_api_usage_patterns_batch(generative_model, tokenizer, library_ids,
                 tokenizer,
                 input_generative_model,
                 num_beams,
-                num_return_sequences
+                num_return_sequences,
+                max_length
             )
 
             temp = input_generative_model.split("[SEP]")
@@ -268,6 +298,7 @@ def generate_api_usage_patterns_batch(generative_model, tokenizer, library_ids,
             constructor = temp[1].strip()
 
             assert(constructor not in temp_dict.get('usage_patterns'))
+            api_usage_patterns = add_params(api_usage_patterns, db_params, id_)
             temp_dict['usage_patterns'][constructor] = api_usage_patterns
         
         assert(temp_dict.get('library_name')==None)
@@ -392,9 +423,10 @@ def initialize_all_components(config):
     classifier_head: a random forest model
     '''
     # load db
-    db_metadata, db_constructor = load_db(
+    db_metadata, db_constructor, db_params = load_db(
         config.get('db_metadata_path'), 
-        config.get('db_constructor_path')
+        config.get('db_constructor_path'),
+        config.get('db_params_path')
     )
 
     # load model
@@ -411,14 +443,14 @@ def initialize_all_components(config):
         config.get('classifier_head_path')
     )
 
-    return db_metadata, db_constructor, model_retrieval, model_generative, tokenizer_generative, model_classifier, classifier_head, tokenizer_classifier
+    return db_metadata, db_constructor, db_params, 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_metadata, db_constructor, db_params,
     config):
     '''
     Function to retrieve relevant libraries, generate API usage patterns, and predict the hw configs
@@ -435,20 +467,19 @@ def make_predictions(input_query,
     Returns:
     predictions (list): a list of dictionary containing the prediction details
     '''
-    print("retrieve library")
     library_ids, library_names = retrieve_libraries(model_retrieval, input_query, db_metadata)
-    
-    print("generate hw patterns")
+
     predictions = generate_api_usage_patterns_batch(
         model_generative,
         tokenizer_generative,
         library_ids,
         db_constructor,
+        db_params,
         config.get('num_beams'),
-        config.get('num_return_sequences')
+        config.get('num_return_sequences'),
+        config.get('max_length_generate')
     )
     
-    print("generate hw config")
     hw_configs = predict_hw_config(
         model_classifier,
         tokenizer_classifier,