fixes

utils/similarity/__init__.py

@@ -7,15 +7,6 @@ except ImportError:
     fastText = None
 
 if fastText is None:
-    print("In order to use slither-simil, you need to install fastText 0.2.0:")
-    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user")
+    print("ERROR: in order to use slither-simil, you need to install fastText 0.2.0:")
+    print("$ pip3 install https://github.com/facebookresearch/fastText/archive/0.2.0.zip --user\n")
     sys.exit(-1)
-
-try:
-    import sklearn
-except ImportError:
-    sklearn = None
-
-if sklearn is None:
-    print("In order to use plot mode in slither-simil, you need to install sklearn:")
-    print("$ pip3 install sklearn --user")

utils/similarity/plot.py

@@ -3,19 +3,30 @@ import sys
 import traceback
 import operator
 import numpy as np
+import tqdm
 import random
 
-from sklearn import decomposition
-import matplotlib.pyplot as plt
+try:
+    from sklearn import decomposition
+    import matplotlib.pyplot as plt
+except ImportError:
+    decomposition = None
+    plt = None
     
 from fastText import load_model
 from .cache import load_cache
 
-logger = logging.getLogger("crytic-pred")
+logger = logging.getLogger("Slither-simil")
 
 def plot(args):
 
+    if decomposition is None or plt is None:
+        print("ERROR: In order to use plot mode in slither-simil, you need to install sklearn and matplotlib:")
+        print("$ pip3 install sklearn matplotlib --user")
+        sys.exit(-1)
+
     try:
+
         model = args.model
         model = load_model(model)
         filename = args.filename
@@ -29,32 +40,32 @@ def plot(args):
             logger.error('The plot mode requieres contract, fname and input parameters.')
             sys.exit(-1)
 
+        logger.info('Loading data..')
         cache = load_cache(infile, model, ext=ext, solc=solc)
-        #save_cache("cache.npz", cache)
 
         data = list()
         fs = list()
+
+        logger.info('Procesing data..')
         for (f,c,n),y in cache.items():
             if c == contract and n == fname:
                 fs.append(f)
                 data.append(y)
-            #r[x] = similarity(fvector, y)
-
        
         data = np.array(data)
         pca = decomposition.PCA(n_components=2)
         tdata = pca.fit_transform(data)
+
+        logger.info('Plotting data..')
         plt.figure()
         assert(len(tdata) == len(fs))
         for ([x,y],l) in zip(tdata, fs):
             x = random.gauss(0, 0.01) + x
             y = random.gauss(0, 0.01) + y
             plt.scatter(x, y, c='blue')
-            plt.text(x-0.001,y+0.001, l.split("_")[1].replace(".sol.ast.compact.json",""))
+            #plt.text(x-0.001,y+0.001, l.split("_")[1].replace(".sol.ast.compact.json",""))
 
-        plt.show()
-        #r = sorted(r.items(), key=operator.itemgetter(1), reverse=True)
-        #for x,score in r[:10]:
+        plt.savefig('plot.png', bbox_inches='tight')
  
     except Exception:
         logger.error('Error in %s' % args.filename)