Update serialization.

Author: qiuosier

ads/llm/__init__.py

@@ -9,4 +9,4 @@
 from ads.llm.langchain.plugins.llm_md import ModelDeploymentTGI
 from ads.llm.langchain.plugins.llm_md import ModelDeploymentVLLM
 from ads.llm.langchain.plugins.embeddings import GenerativeAIEmbeddings
-from ads.llm.load import load
+from ads.llm.serialize import load

ads/llm/load.py renamed to ads/llm/serialize.py

@@ -4,8 +4,10 @@
 # Copyright (c) 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import json
 from typing import Any, Dict, List, Optional
 from langchain.load.load import load as lc_load
+from langchain.load.serializable import Serializable
 
 
 def load(
@@ -14,10 +16,10 @@ def load(
     secrets_map: Optional[Dict[str, str]] = None,
     valid_namespaces: Optional[List[str]] = None,
 ) -> Any:
-    """Revive a LangChain class from a JSON object. Use this if you already
+    """Revive an ADS/LangChain class from a JSON object. Use this if you already
     have a parsed JSON object, eg. from `json.load` or `json.loads`.
 
-    This is a drop in replacement for load() in langchain.load.load to support loading ADS LangChain compatible component.
+    This is a drop in replacement for load() in langchain.load.load to support loading compatible class from ADS.
 
     Args:
         obj: The object to load.
@@ -33,3 +35,35 @@ def load(
     if "ads" not in valid_namespaces:
         valid_namespaces.append("ads")
     return lc_load(obj, secrets_map=secrets_map, valid_namespaces=valid_namespaces)
+
+
+def default(obj: Any) -> Any:
+    """Calls the to_json() method to serialize the object.
+
+    Parameters
+    ----------
+    obj : Any
+        The object to be serialized.
+
+    Returns
+    -------
+    Any
+        The serialized representation of the object.
+
+    Raises
+    ------
+    TypeError
+        If the object is not LangChain serializable.
+    """
+    if isinstance(obj, Serializable) and obj.is_lc_serializable():
+        return obj.to_json()
+    raise TypeError(f"Serialization of {type(obj)} is not supported.")
+
+
+def dump(obj: Any) -> Dict[str, Any]:
+    """Return a json dict representation of an object.
+
+    This is a drop in replacement of the dumpd() in langchain.load.dump,
+    except that this method will raise TypeError when the object is not serializable.
+    """
+    return json.loads(json.dumps(obj, default=default))

tests/unitary/with_extras/langchain/test_guardrails.py

@@ -11,11 +11,10 @@
 from langchain.llms.base import LLM
 from langchain.prompts import PromptTemplate
 from langchain.schema.runnable import RunnableMap, RunnablePassthrough
-from langchain.load import dumpd
 from ads.llm.guardrails import HuggingFaceEvaluation
 from ads.llm.guardrails.base import BlockedByGuardrail, GuardrailIO
 from ads.llm.chain import GuardrailSequence
-from ads.llm.load import load
+from ads.llm.serialize import load, dump
 
 
 class FakeLLM(LLM):
@@ -64,7 +63,7 @@ def test_toxicity_without_threshold(self):
         chain = self.FAKE_LLM | toxicity
         output = chain.invoke(self.TOXIC_CONTENT)
         self.assertEqual(output, self.TOXIC_CONTENT)
-        serialized = dumpd(chain)
+        serialized = dump(chain)
         chain = load(serialized, valid_namespaces=["tests"])
         output = chain.invoke(self.TOXIC_CONTENT)
         self.assertEqual(output, self.TOXIC_CONTENT)
@@ -77,7 +76,7 @@ def test_toxicity_with_threshold(self):
         chain = self.FAKE_LLM | toxicity
         with self.assertRaises(BlockedByGuardrail):
             chain.invoke(self.TOXIC_CONTENT)
-        serialized = dumpd(chain)
+        serialized = dump(chain)
         chain = load(serialized, valid_namespaces=["tests"])
         with self.assertRaises(BlockedByGuardrail):
             chain.invoke(self.TOXIC_CONTENT)
@@ -94,7 +93,7 @@ def test_toxicity_without_exception(self):
         chain = self.FAKE_LLM | toxicity
         output = chain.invoke(self.TOXIC_CONTENT)
         self.assertEqual(output, toxicity.custom_msg)
-        serialized = dumpd(chain)
+        serialized = dump(chain)
         chain = load(serialized, valid_namespaces=["tests"])
         output = chain.invoke(self.TOXIC_CONTENT)
         self.assertEqual(output, toxicity.custom_msg)
@@ -109,7 +108,7 @@ def test_toxicity_return_metrics(self):
         self.assertIsInstance(output, dict)
         self.assertEqual(output["output"], self.TOXIC_CONTENT)
         self.assertGreater(output["metrics"]["toxicity"][0], 0.2)
-        serialized = dumpd(chain)
+        serialized = dump(chain)
         chain = load(serialized, valid_namespaces=["tests"])
         output = chain.invoke(self.TOXIC_CONTENT)
         self.assertIsInstance(output, dict)
@@ -123,9 +122,11 @@ class GuardrailSequenceTests(GuardrailTestsBase):
     def test_guardrail_sequence_with_template_and_toxicity(self):
         template = PromptTemplate.from_template("Tell me a joke about {subject}")
         map_input = RunnableMap(subject=RunnablePassthrough())
-        toxicity = HuggingFaceEvaluation(path="toxicity", load_args=self.LOAD_ARGS)
+        toxicity = HuggingFaceEvaluation(
+            path="toxicity", load_args=self.LOAD_ARGS, select=min
+        )
         chain = GuardrailSequence.from_sequence(
             map_input | template | self.FAKE_LLM | toxicity
         )
-        output = chain.run("cats")
+        output = chain.run("cats", num_generations=5)
         self.assertIsInstance(output, GuardrailIO)