docs: move and edit existing linear-regression tutorial snippet (#991)

* add new alias '__version__' * remove accidental changes * correct assignment * docs: edit existing code snippet for linear regression and add to /snippets * return type * remove filtering * input data updated * remove score and predict Co-authored-by: Tim Sweña (Swast) <[email protected]> * delete blank line Co-authored-by: Tim Sweña (Swast) <[email protected]> * remove test data Co-authored-by: Tim Sweña (Swast) <[email protected]> * remove test checks Co-authored-by: Tim Sweña (Swast) <[email protected]> * remove test check Co-authored-by: Tim Sweña (Swast) <[email protected]> * add variable input to allow random string to be passed and saved as the model Co-authored-by: Tim Sweña (Swast) <[email protected]> * add line that writes the dataframe to big query * update comment with the correct tutorial model reference * update comment with the correct tutorial model reference * remove extra line --------- Co-authored-by: Tim Sweña (Swast) <[email protected]>
googleapis · Sep 24, 2024 · 4cb62fd · 4cb62fd
1 parent 208a984
commit 4cb62fd
Showing 1 changed file with 42 additions and 0 deletions.
diff --git a/samples/snippets/linear_regression_tutorial_test.py b/samples/snippets/linear_regression_tutorial_test.py
@@ -0,0 +1,42 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def test_linear_regression(random_model_id: str) -> None:
+    your_model_id = random_model_id
+    # [START bigquery_dataframes_bqml_linear_regression]
+    from bigframes.ml.linear_model import LinearRegression
+    import bigframes.pandas as bpd
+
+    # Load data from BigQuery
+    bq_df = bpd.read_gbq("bigquery-public-data.ml_datasets.penguins")
+
+    # Drop rows with nulls to get training data
+    training_data = bq_df.dropna(subset=["body_mass_g"])
+
+    # Specify your feature (or input) columns and the label (or output) column:
+    feature_columns = training_data.drop(columns=["body_mass_g"])
+    label_columns = training_data[["body_mass_g"]]
+
+    # Create the linear model
+    model = LinearRegression()
+    model.fit(feature_columns, label_columns)
+    model.to_gbq(
+        your_model_id,  # For example: "bqml_tutorial.penguins_model"
+        replace=True,
+    )
+    # [END bigquery_dataframes_bqml_linear_regression]
+    assert feature_columns is not None
+    assert label_columns is not None
+    assert model is not None