[翻訳] LangChainによるDatabricksでサービングされているLLMの呼び出し

[Databricks \| 🦜️🔗 Langchain](https://python.langchain.com/docs/modules/model_io/models/llms/integrations/databricks.html)の翻訳です。

:::note warn
本書は抄訳であり内容の正確性を保証するものではありません。正確な内容に関しては原文を参照ください。
:::

[Databricks](https://www.databricks.com/jp)レイクハウスプラットフォームは一つのプラットフォームでデータ、分析、AIを統合します。

このサンプルノートブックでは、LangChainのLLMとしてどのようにDatabricksエンドポイントをラッピングするのかを説明します。2つタイプのエンドポイントがサポートされています。

- サービングエンドポイント、プロダクションや開発で推奨。
- クラスタードライバーのプロキシーアプリ、インタラクティブな開発で推奨。

```py
from langchain.llms import Databricks
```

# サービングエンドポイントのラッピング

前提条件

- [Databricksサービングエンドポイント](https://docs.databricks.com/machine-learning/model-serving/index.html)にLLMが登録、デプロイされている。
- エンドポイントへの[Can Query権限](https://docs.databricks.com/security/auth-authz/access-control/serving-endpoint-acl.html)がある。

期待されるMLflowモデルのシグネチャは: 

- 入力: `[{"name": "prompt", "type": "string"}, {"name": "stop", "type": "list[string]"}]`
- 出力: `[{"type": "string"}]`

モデルのシグネチャに互換性がない、あるいは追加の設定を挿入したい場合には、それぞれ、`transform_input_fn`と`transform_output_fn`を設定することができます。

```py
# If running a Databricks notebook attached to an interactive cluster in "single user"
# or "no isolation shared" mode, you only need to specify the endpoint name to create
# a `Databricks` instance to query a serving endpoint in the same workspace.
llm = Databricks(endpoint_name="dolly")

llm("How are you?")
```
```
    'I am happy to hear that you are in good health and as always, you are appreciated.'
```

```py
llm("How are you?", stop=["."])
```
```
    'Good'
```

```py
# Otherwise, you can manually specify the Databricks workspace hostname and personal access token
# or set `DATABRICKS_HOST` and `DATABRICKS_TOKEN` environment variables, respectively.
# See https://docs.databricks.com/dev-tools/auth.html#databricks-personal-access-tokens
# We strongly recommend not exposing the API token explicitly inside a notebook.
# You can use Databricks secret manager to store your API token securely.
# See https://docs.databricks.com/dev-tools/databricks-utils.html#secrets-utility-dbutilssecrets

import os

os.environ["DATABRICKS_TOKEN"] = dbutils.secrets.get("myworkspace", "api_token")

llm = Databricks(host="myworkspace.cloud.databricks.com", endpoint_name="dolly")

llm("How are you?")
```
```
    'I am fine. Thank you!'
```

```py
# If the serving endpoint accepts extra parameters like `temperature`,
# you can set them in `model_kwargs`.
llm = Databricks(endpoint_name="dolly", model_kwargs={"temperature": 0.1})

llm("How are you?")
```
```
    'I am fine.'
```

```py
# Use `transform_input_fn` and `transform_output_fn` if the serving endpoint
# expects a different input schema and does not return a JSON string,
# respectively, or you want to apply a prompt template on top.


def transform_input(**request):
    full_prompt = f"""{request["prompt"]}
    Be Concise.
    """
    request["prompt"] = full_prompt
    return request


llm = Databricks(endpoint_name="dolly", transform_input_fn=transform_input)

llm("How are you?")
```
```
    'I’m Excellent. You?'
```

# クラスタードライバーのプロキシーアプリのラッピング

前提条件

- "single user"あるいは"no isolation shared"モードのDatabricksインタラクティブクラスターにLLMがロードされている。
- JSON入力/出力を行うHTTP POSTを用いて、モデルを`"/"`でサーブするために、ドライバーノードでローカルのHTTPサーバーが稼働している。
- `[3000, 8000]`の間のポート番号を使用し、ドライバーのIPアドレスをリッスンするか、ローカルホストのみではなくシンプルに`0.0.0.0`をリッスンする。
- クラスターへの"Can Attach To"権限がある。

期待される(JSONスキーマを用いた)サーバースキーマは:

- 入力: 

```json
{"type": "object",
 "properties": {
    "prompt": {"type": "string"},
     "stop": {"type": "array", "items": {"type": "string"}}},
  "required": ["prompt"]}
```

- 出力: `{"type": "string"}`

サーバーのシグネチャに互換性がない、あるいは追加の設定を挿入したい場合には、それぞれ、`transform_input_fn`と`transform_output_fn`を設定することができます。

LLMをサーブするためにドライバープロキシーアプリを実行する最小限のサンプルを示します。

```py
from flask import Flask, request, jsonify
import torch
from transformers import pipeline, AutoTokenizer, StoppingCriteria

model = "databricks/dolly-v2-3b"
tokenizer = AutoTokenizer.from_pretrained(model, padding_side="left")
dolly = pipeline(model=model, tokenizer=tokenizer, trust_remote_code=True, device_map="auto")
device = dolly.device

class CheckStop(StoppingCriteria):
    def __init__(self, stop=None):
        super().__init__()
        self.stop = stop or []
        self.matched = ""
        self.stop_ids = [tokenizer.encode(s, return_tensors='pt').to(device) for s in self.stop]
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
        for i, s in enumerate(self.stop_ids):
            if torch.all((s == input_ids[0][-s.shape[1]:])).item():
                self.matched = self.stop[i]
                return True
        return False

def llm(prompt, stop=None, **kwargs):
  check_stop = CheckStop(stop)
  result = dolly(prompt, stopping_criteria=[check_stop], **kwargs)
  return result[0]["generated_text"].rstrip(check_stop.matched)

app = Flask("dolly")

@app.route('/', methods=['POST'])
def serve_llm():
  resp = llm(**request.json)
  return jsonify(resp)

app.run(host="0.0.0.0", port="7777")
```

サーバーが稼働すると、LLMとしてラッピングする`Databricks`インスタンスを作成することができます。

```py
# If running a Databricks notebook attached to the same cluster that runs the app,
# you only need to specify the driver port to create a `Databricks` instance.
llm = Databricks(cluster_driver_port="7777")

llm("How are you?")
```
```
    'Hello, thank you for asking. It is wonderful to hear that you are well.'
```

```py
# Otherwise, you can manually specify the cluster ID to use,
# as well as Databricks workspace hostname and personal access token.

llm = Databricks(cluster_id="0000-000000-xxxxxxxx", cluster_driver_port="7777")

llm("How are you?")
```
```
    'I am well. You?'
```

```py
# If the app accepts extra parameters like `temperature`,
# you can set them in `model_kwargs`.
llm = Databricks(cluster_driver_port="7777", model_kwargs={"temperature": 0.1})

llm("How are you?")
```
```
    'I am very well. It is a pleasure to meet you.'
```

```py
# Use `transform_input_fn` and `transform_output_fn` if the app
# expects a different input schema and does not return a JSON string,
# respectively, or you want to apply a prompt template on top.


def transform_input(**request):
    full_prompt = f"""{request["prompt"]}
    Be Concise.
    """
    request["prompt"] = full_prompt
    return request


def transform_output(response):
    return response.upper()


llm = Databricks(
    cluster_driver_port="7777",
    transform_input_fn=transform_input,
    transform_output_fn=transform_output,
)

llm("How are you?")
```
```
    'I AM DOING GREAT THANK YOU.'
```