【Google Cloud イベント体験記】LangChain を用いた対話型生成 AI アプリの作り方

! pip install google-cloud-aiplatform \
    langchain==0.0.323 \
    chromadb==0.3.26 \
    pydantic==1.10.8 \
    typing-inspect==0.8.0 \
    typing_extensions==4.5.0 \
    pandas \
    datasets \
    google-api-python-client \
    pypdf \
    faiss-cpu \
    transformers \
    config \
    --upgrade --user

import IPython

app = IPython.Application.instance()
app.kernel.do_shutdown(True)

import time

import langchain
from google.cloud import aiplatform
from langchain.chat_models import ChatVertexAI
from langchain.embeddings import VertexAIEmbeddings
from langchain.llms import VertexAI
from langchain.schema import HumanMessage, SystemMessage
from pydantic import BaseModel
from langchain.prompts import PromptTemplate

print(f"LangChain version: {langchain.__version__}")
print(f"Vertex AI SDK version: {aiplatform.__version__}")

LangChain version: 0.0.323
Vertex AI SDK version: 1.67.1

# Embeddings APIのレート制限のためのユーティリティ関数
def rate_limit(max_per_minute: int) -> None:
    period = 60 / max_per_minute
    print("Waiting")
    while True:
        before = time.time()
        yield
        after = time.time()
        elapsed = after - before
        sleep_time = max(0, period - elapsed)
        if sleep_time > 0:
            print(".", end="")
            time.sleep(sleep_time)

class CustomVertexAIEmbeddings(VertexAIEmbeddings, BaseModel):
    requests_per_minute: int
    num_instances_per_batch: int

    # embed_documents メソッドのオーバーライド
    def embed_documents(self, texts: list) -> list:
        limiter = rate_limit(self.requests_per_minute)
        results = []
        docs = list(texts)

        while docs:
            head, docs = (
                docs[: self.num_instances_per_batch],
                docs[self.num_instances_per_batch :],
            )
            chunk = self.client.get_embeddings(head)
            results.extend(chunk)
            next(limiter)

        return [r.to_numpy() for r in results]

llm = VertexAI(
    model_name="text-bison",
    max_output_tokens=256,
    temperature=0.1,
    top_p=0.8,
    top_k=40,
    verbose=True,
)

EMBEDDING_QPM = 100
EMBEDDING_NUM_BATCH = 5

embeddings = CustomVertexAIEmbeddings(
    requests_per_minute=EMBEDDING_QPM,
    num_instances_per_batch=EMBEDDING_NUM_BATCH,
)

example_prompt = PromptTemplate(
    input_variables=["input", "output"],
    template="Example Input: {input}\nExample Output: {output}",
)

examples = [
    {"input": "pirate", "output": "ship"},
    {"input": "pilot", "output": "plane"},
    {"input": "driver", "output": "car"},
    {"input": "tree", "output": "ground"},
    {"input": "bird", "output": "nest"},
]

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    embeddings,
    FAISS,
    k=2,
)
similar_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="指定したものが通常置かれている場所を示す",
    suffix="Input: {noun}\nOutput:",
    input_variables=["noun"],
)
my_noun = "student"

print(similar_prompt.format(noun=my_noun))
print(llm(similar_prompt.format(noun=my_noun)))

Waiting
指定したものが通常置かれている場所を示す

Example Input: driver
Example Output: car

Example Input: pilot
Example Output: plane

Input: student
Output:
 school