V1beta3 (#67)

* Working on v1beta3

Add support for batched embeddings
Update to v0.2
Add get_tuned_model, get_base_model, get_model.
Add list_models, list_tuned_models.
Add create_tuned_model.
Support batched embeddings.

Change-Id: I917576a61719fcd90cca581506e15826b67cae88

* Implement {update,delete}_tuned_model.

Change-Id: Idd276934666f7fa8da034e59a80e4e6047cea4d5

* Decode TunedModel to python objects.

Change-Id: I4376012686b53065f9372bd8dd6a340d3e0ca43a

* Add more model service tests + bug fixes + pytype.

Change-Id: Iaa1b4e75b5c075d1b0b6b383d37da8d254b0715b

* Upgrade the Operation returned by create_tuned_model.

Updated into colab here: https://colab.sandbox.google.com/drive/1Eros5eLfAl4pNfsQ2sftyx5khq2TrK_V?resourcekey=0-ByOp_F6lxWCuKVJM_tUjeQ

Change-Id: Ic3e40f25abc6eba69a8866fe855c7ecb83c4c4c3

* Implement the operations service.

Change-Id: I29e92a4c8697f165deb9782599b21d98eb869737

* fix update_tuned_model

Change-Id: Iaacaaa08c1807e13ad520804dd20eb1e54a9b34d

* Remove decoding error workaround + format + use the 0.3.2 v1beta3 package.

Change-Id: I6e329c635aa51e67c50b2318a506455c21c8d988

* Fix tests & remove pytz

Fix tuned_model/base_model/source_model handling so nested tuned models work
Remove the last 'v1main'

Change-Id: Ia07c169d0bfb0782edd0115b5d0418c9eb3eb040

* black .

* fix test

* add tqdm requirement

* Use union to fix py3.9

* Use union to fix py3.9

* black .

Author: MarkDaoust

RELEASE.md

@@ -0,0 +1,16 @@
+
+## v0.2
+
+- More flexible safety settings: 
+  - accept strings, ints for enums.
+  - accept a `{category:threshold}` dict in addition to the
+    list of dicts format (`[{"category": category, "threshold": threshold}, ...]`).
+- Add support for batched embeddings.
+- Add support for tuning:
+  - Add `get_{base,tuned}_model`.
+  - Add `list_tuned_models`.
+  - Add `create_tuned_model`.
+
+## v0.1
+
+Initial version

docs/build_docs.py

@@ -29,15 +29,15 @@
 from absl import flags
 
 import google
-from google import generativeai
-from google.ai import generativelanguage
+from google import generativeai as palm
+from google.ai import generativelanguage_v1beta3 as glm
 
 from tensorflow_docs.api_generator import generate_lib
 from tensorflow_docs.api_generator import public_api
 
 import yaml
 
-google.ai.generativelanguage.__doc__ = """\
+glm.__doc__ = """\
 This package, `google.ai.generativelanguage`, is a low-level auto-generated client library for the PaLM API.
 
 ```posix-terminal
@@ -54,7 +54,7 @@
 when initializing a client:
 
 ```
-from google.ai import generativelanguage as glm
+from google.ai import generativelanguage_v1beta3 as glm
 
 client = glm.DiscussServiceClient(
     client_options={'api_key':'YOUR_API_KEY'})
@@ -186,8 +186,8 @@ def gen_api_docs():
         py_modules=[("google", google)],
         # Replace `tensorflow_docs.api_generator` with your module, here.
         base_dir=(
-            pathlib.Path(google.generativeai.__file__).parent,
-            pathlib.Path(google.ai.generativelanguage.__file__).parent.parent,
+            pathlib.Path(palm.__file__).parent,
+            pathlib.Path(glm.__file__).parent.parent,
         ),
         code_url_prefix=(
             _CODE_URL_PREFIX.value,

google/generativeai/__init__.py

@@ -79,7 +79,15 @@
 from google.generativeai.text import generate_embeddings
 
 from google.generativeai.models import list_models
+from google.generativeai.models import list_tuned_models
+
 from google.generativeai.models import get_model
+from google.generativeai.models import get_base_model
+from google.generativeai.models import get_tuned_model
+
+from google.generativeai.models import create_tuned_model
+from google.generativeai.models import update_tuned_model
+from google.generativeai.models import delete_tuned_model
 
 from google.generativeai.client import configure
 

google/generativeai/client.py

@@ -17,11 +17,13 @@
 import os
 from typing import cast, Optional, Union
 
-import google.ai.generativelanguage as glm
+import google.ai.generativelanguage_v1beta3 as glm
 
 from google.auth import credentials as ga_credentials
 from google.api_core import client_options as client_options_lib
 from google.api_core import gapic_v1
+from google.api_core import operations_v1
+
 from google.generativeai import version
 
 
@@ -32,6 +34,7 @@
 default_discuss_async_client = None
 default_model_client = None
 default_text_client = None
+default_operations_client = None
 
 
 def configure(
@@ -64,6 +67,7 @@ def configure(
     global default_discuss_client
     global default_model_client
     global default_text_client
+    global default_operations_client
 
     if isinstance(client_options, dict):
         client_options = client_options_lib.from_dict(client_options)
@@ -112,6 +116,7 @@ def configure(
     default_discuss_client = None
     default_text_client = None
     default_model_client = None
+    default_operations_client = None
 
 
 def get_default_discuss_client() -> glm.DiscussServiceClient:
@@ -125,7 +130,7 @@ def get_default_discuss_client() -> glm.DiscussServiceClient:
     return default_discuss_client
 
 
-def get_default_text_client():
+def get_default_text_client() -> glm.TextServiceClient:
     global default_text_client
     if default_text_client is None:
         # Attempt to configure using defaults.
@@ -149,7 +154,7 @@ def get_default_discuss_async_client() -> glm.DiscussServiceAsyncClient:
     return default_discuss_async_client
 
 
-def get_default_model_client():
+def get_default_model_client() -> glm.ModelServiceClient:
     global default_model_client
     if default_model_client is None:
         # Attempt to configure using defaults.
@@ -158,3 +163,12 @@ def get_default_model_client():
         default_model_client = glm.ModelServiceClient(**default_client_config)
 
     return default_model_client
+
+
+def get_default_operations_client() -> operations_v1.OperationsClient:
+    global default_operations_client
+    if default_operations_client is None:
+        model_client = get_default_model_client()
+        default_operations_client = model_client._transport.operations_client
+
+    return default_operations_client

google/generativeai/discuss.py

@@ -20,7 +20,7 @@
 
 from typing import Iterable, List, Optional, Union
 
-import google.ai.generativelanguage as glm
+import google.ai.generativelanguage_v1beta3 as glm
 
 from google.generativeai.client import get_default_discuss_client
 from google.generativeai.client import get_default_discuss_async_client
@@ -271,7 +271,7 @@ def _make_message_prompt(
 
 def _make_generate_message_request(
     *,
-    model: model_types.ModelNameOptions | None,
+    model: model_types.AnyModelNameOptions | None,
     context: str | None = None,
     examples: discuss_types.ExamplesOptions | None = None,
     messages: discuss_types.MessagesOptions | None = None,
@@ -313,7 +313,7 @@ def inner(f):
 
 def chat(
     *,
-    model: model_types.ModelNameOptions | None = "models/chat-bison-001",
+    model: model_types.AnyModelNameOptions | None = "models/chat-bison-001",
     context: str | None = None,
     examples: discuss_types.ExamplesOptions | None = None,
     messages: discuss_types.MessagesOptions | None = None,
@@ -411,7 +411,7 @@ def chat(
 @set_doc(chat.__doc__)
 async def chat_async(
     *,
-    model: model_types.ModelNameOptions | None = "models/chat-bison-001",
+    model: model_types.AnyModelNameOptions | None = "models/chat-bison-001",
     context: str | None = None,
     examples: discuss_types.ExamplesOptions | None = None,
     messages: discuss_types.MessagesOptions | None = None,
@@ -567,7 +567,7 @@ def count_message_tokens(
     context: str | None = None,
     examples: discuss_types.ExamplesOptions | None = None,
     messages: discuss_types.MessagesOptions | None = None,
-    model: model_types.ModelNameOptions = DEFAULT_DISCUSS_MODEL,
+    model: model_types.AnyModelNameOptions = DEFAULT_DISCUSS_MODEL,
     client: glm.DiscussServiceAsyncClient | None = None,
 ):
     model = model_types.make_model_name(model)