Improve description and figure

app.py

@@ -5,12 +5,15 @@ import numpy as np
 from sklearn import datasets, linear_model
 from sklearn.metrics import mean_squared_error, r2_score
 
+import matplotlib
+matplotlib.use('agg')
+
 FIGSIZE = (10,10)
 
-feature_names = ["age", "body-mass index (BMI)", "blood pressure", 
-                "total serum cholesterol", "low-density lipoproteins (LDL)", 
-                "high-density lipoproteins (HDL)", "total cholesterol / HDL", 
-                "log of serum triglycerides level (possibly)","blood sugar level"]
+feature_names = ["Age", "Body-Mass Index (BMI)", "Blood Pressure", 
+                "Total serum Cholesterol", "Low-Density Lipoproteins (LDL)", 
+                "High-Density Lipoproteins (HDL)", "Total cholesterol / HDL", 
+                "log(Serum Triglycerides Level) (possibly)","Blood Sugar Level"]
 
 def create_dataset(feature_id=2):
     # Load the diabetes dataset
@@ -59,15 +62,24 @@ def train_model(input_data):
     plt.scatter(diabetes_X_test, diabetes_y_test, color="black")
     plt.plot(diabetes_X_test, diabetes_y_pred, color="blue", linewidth=3)
 
+    plt.xlabel(input_data, fontsize=18)
+    plt.ylabel("Disease progression", fontsize=18)
+
     plt.xticks(())
     plt.yticks(())
 
-
-
     return fig, regr.coef_, mse, r2
 
 title = "Linear Regression Example 📈"
-description = "The example shows how linear regression attempts to draw a straight line that will best minimize the residual sum of squares between the observed responses in the dataset"
+description = """The example shows how linear regression attempts to draw a straight line that will best minimize the residual sum of squares between the observed responses in the dataset. 
+
+                The diabetes dataset contains baseline variables (features), age, sex, body mass index, average blood pressure, and six blood serum measurements that were obtained for 442 diabetes patients.
+                The predictive variable is a quantitative measure of the disease progression one year after the baseline.
+
+                When selecting a feature from the drop-down menu, a linear regression model is trained for the specific feature and the predictive variable.
+                The figure shows a scatter plot of the test set as well as the linear model (line).
+                The mean square error and R2 scores are calculated using the test set and they are printed, along with the regression coefficiet of the model. 
+                """
 with gr.Blocks() as demo:
     gr.Markdown(f"## {title}")
     gr.Markdown(description)
@@ -79,8 +91,8 @@ with gr.Blocks() as demo:
             with gr.Column():
                 input_data = gr.Dropdown(choices=feature_names, label="Feature", value="body-mass index")
                 coef = gr.Textbox(label="Coefficients")
-                mse = gr.Textbox(label="MSE")
-                r2 = gr.Textbox(label="R2")
+                mse = gr.Textbox(label="Mean Squared Error (MSE)")
+                r2 = gr.Textbox(label="R2 score")
 
         input_data.change(fn=train_model, inputs=[input_data], outputs=[plot, coef, mse, r2], queue=False)