fixes + enable contract sampling during training

utils/similarity/__main__.py

@@ -51,6 +51,12 @@ def parse_args():
                         dest='fname',
                         help='Function name')
 
+    parser.add_argument('--nsamples',
+                        action='store',
+                        type=int,
+                        dest='nsamples',
+                        help='Number of contract samples used for training')
+
     parser.add_argument('--input',
                         action='store',
                         dest='input',

utils/similarity/encode.py

@@ -49,8 +49,8 @@ def ntype(_type):
     else:
         _type = str(_type)
 
-    _type = _type.replace("_memory","")
-    _type = _type.replace("_storage_ref","")
+    _type = _type.replace(" memory","")
+    _type = _type.replace(" storage ref","")
 
     if "struct" in _type:
         return "struct"

utils/similarity/train.py

@@ -12,23 +12,26 @@ logger = logging.getLogger("Slither-simil")
 def train(args):
 
     try:
+        last_data_train_filename = "last_data_train.txt"
         model_filename = args.model
         solc = args.solc
         dirname = args.input
         ext = args.filter
+        nsamples = args.nsamples
 
         if dirname is None:
             logger.error('The train mode requires the input parameter.')
             sys.exit(-1)
 
-        contracts = load_contracts(dirname, ext=ext, nsamples=None)
-        with open("data.txt", 'w') as f:
+        contracts = load_contracts(dirname, ext=ext, nsamples=nsamples)
+        logger.info('Saving extracted data into', last_data_train_filename)
+        with open(last_data_train_filename, 'w') as f:
             for contract in contracts: 
                 for function,ir in encode_contract(contract,solc).items():
                     if ir != []:
                         f.write(" ".join(ir)+"\n")
     
-        model = train_unsupervised(input='data.txt', model='skipgram')
+        model = train_unsupervised(input=last_data_train_filename, model='skipgram')
         model.save_model(model_filename)
         print(model.get_words())