Update mymodel.py

mymodel.py

@@ -12,6 +12,8 @@ from sklearn.tree import DecisionTreeClassifier
 from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
 import seaborn as sn
 import matplotlib.pyplot as plt
+from concrete.ml.sklearn import DecisionTreeClassifier as ConcreteDecisionTreeClassifier
+from concrete.ml.sklearn import XGBClassifier as ConcreteXGBClassifier
 
 
 class DiseasePrediction:
@@ -100,11 +102,9 @@ class DiseasePrediction:
         if self.model_name == 'mnb':
             self.clf = MultinomialNB()
         elif self.model_name == 'decision_tree':
-            self.clf = DecisionTreeClassifier(criterion=self.config['model']['decision_tree']['criterion'])
-        elif self.model_name == 'random_forest':
-            self.clf = RandomForestClassifier(n_estimators=self.config['model']['random_forest']['n_estimators'])
+            self.clf = ConcreteDecisionTreeClassifier(criterion=self.config['model']['decision_tree']['criterion'])
         elif self.model_name == 'gradient_boost':
-            self.clf = GradientBoostingClassifier(n_estimators=self.config['model']['gradient_boost']['n_estimators'],
+            self.clf = ConcreteXGBClassifier(n_estimators=self.config['model']['gradient_boost']['n_estimators'],
                                                   criterion=self.config['model']['gradient_boost']['criterion'])
         return self.clf