OpenAI Embedding APIを使ってみる【Python】

import pandas as pd
import tiktoken
from openai import OpenAI

embedding_model = "text-embedding-3-small"
embedding_encoding = "cl100k_base"
max_tokens = 1500  # embedding のパラメータ設定

# 表にまとめた文章を読み込み、カラム名を「title」と「text」に変更
df = pd.read_csv("（表にまとめた文章のファイル名）.csv")
df.columns = ['title', 'text']

# textのトークン数を計算し、新しい列n_tokensに格納
tokenizer = tiktoken.get_encoding(embedding_encoding)
df['n_tokens'] = df.text.apply(lambda x: len(tokenizer.encode(x)))


def split_into_many(text, max_tokens = 500):
    '''
    テキストを最大トークン数に分割する関数
    '''

    # テキストを文ごとに分割し、各文のトークン数を取得
    sentences = text.split('。')
    n_tokens = [len(tokenizer.encode(" " + sentence)) for sentence in sentences]

    chunks = []
    tokens_so_far = 0
    chunk = []

    # 各文とトークンを組み合わせてループ処理
    for sentence, token in zip(sentences, n_tokens):

        # これまでのトークン数と現在の文のトークン数を合計した値が
        # 最大トークン数を超える場合は、チャンクをチャンクのリストに追加し、
        # チャンクとトークン数をリセット
        if tokens_so_far + token > max_tokens:
            chunks.append(". ".join(chunk) + ".")
            chunk = []
            tokens_so_far = 0

        # 現在の文のトークン数が最大トークン数より大きい場合は、次の文へ進む
        if token > max_tokens:
            continue

        # それ以外の場合は、文をチャンクに追加し、トークン数を合計に追加
        chunk.append(sentence)
        tokens_so_far += token + 1

    # 最後のチャンクをチャンクのリストに追加
    if chunk:
        chunks.append(". ".join(chunk) + ".")
    return chunks


# 短縮されたテキストを格納するためのリスト
shortened = []

# DataFrameの各行に対してループ処理
for row in df.iterrows():
    # テキストがNoneの場合は、次の行へ進む
    if row[1]['text'] is None:
        continue
    # トークン数が最大トークン数より大きい場合は、テキストを
    # 「shortened」リストに追加
    if row[1]['n_tokens'] > max_tokens:
        shortened += split_into_many(row[1]['text'])

    # それ以外の場合は、テキストをそのまま「shortened」リストに追加
    else:
        shortened.append(row[1]['text'])

# 「shortened」をもとに新しいDataFrameを作成し、列名を「text」とする
df = pd.DataFrame(shortened, columns = ['text'])

# 各「text」のトークン数を計算し、新しい列「n_tokens」に格納
df['n_tokens'] = df.text.apply(lambda x: len(tokenizer.encode(x)))


client = OpenAI()

def get_embedding(text, model=embedding_model):
   text = text.replace("\n", " ")
   return client.embeddings.create(input = [text], model=model).data[0].embedding


# 「text」列のテキストに対してembeddingを行い、CSVファイルに保存
df["embeddings"] = df.text.apply(lambda x: get_embedding(x, model=embedding_model))
df.to_csv('./embeddings.csv')

import pandas as pd
import numpy as np
import os
from openai import OpenAI
from scipy.spatial.distance import cosine


def create_context(question, df, max_len=1800):
    """
    質問と学習データを比較して、コンテキストを作成する関数
    """

    client = OpenAI(
        api_key=os.environ.get("OPENAI_API_KEY"),
    )

    # 質問をベクトル化
    model = "text-embedding-3-small"
    q_embeddings = client.embeddings.create(input=question,
                                            model=model).data[0].embedding

    # q_embeddingsを１次配列に変換
    np_array = np.array(q_embeddings)
    flat_array = np_array.flatten()

    # 文字列を数値リストに変換する関数
    def convert_string_to_list(str_emb):
        # 文字列をカンマで分割し、数値に変換
        return [float(num) for num in str_emb.strip('[]').split(',')]

    # embeddings列の各要素を数値リストに変換
    df['embeddings'] = df['embeddings'].apply(convert_string_to_list)



    # 質問と学習データと比較してコサイン類似度を計算し、「distances」という列に類似度を格納
    df['distances'] = df['embeddings'].apply(lambda x: cosine(np.squeeze(flat_array), np.squeeze(np.array(x))))

    # コンテキストを格納するためのリスト
    returns = []
    # コンテキストの現在の長さ
    cur_len = 0

    # 学習データを類似度順にソートし、トークン数の上限までコンテキストに追加する
    for _, row in df.sort_values('distances', ascending=True).iterrows():
        # テキストの長さを現在の長さに加える
        cur_len += row['n_tokens'] + 4

        # テキストが長すぎる場合はループを終了
        if cur_len > max_len:
            break

        # コンテキストのリストにテキストを追加する
        returns.append(row["text"])

    # コンテキストを結合して返す
    return "\n\n###\n\n".join(returns)


def answer_question(question, conversation_history):
    """
    コンテキストに基づいて質問に答える関数
    """

    client = OpenAI(
        api_key=os.environ.get("OPENAI_API_KEY"),
    )

    # 学習データを読み込む
    df = pd.read_csv('embeddings.csv', encoding="utf-8")

    # 質問と学習データを比較してコンテキストを作成
    context = create_context(question, df, max_len=200)
    # プロンプトを作成し、会話の履歴に追加
    prompt = f"あなたはとあるホテルのスタッフです。コンテキストに基づいて、お客様からの質問に丁寧に答えてください。コンテキストが質問に対して回答できない場合は「わかりません」と答えてください。\n\nコンテキスト: {context}\n\n---\n\n質問: {question}\n回答:"
    conversation_history.append({"role": "user", "content": prompt})

    try:
        # ChatGPTからの回答を生成
        response = client.chat.completions.create(
            model = "gpt-3.5-turbo",
            temperature = 0.5,
            messages = conversation_history,
        )

        # ChatGPTからの回答を返す
        return response.choices[0].message.content.strip()

    except Exception as e:
        # エラーが発生した場合は空の文字列を返す
        print(e)
        return ""

import os
from openai import OpenAI
from search import answer_question

client = OpenAI(
        api_key=os.environ.get("OPENAI_API_KEY"),
)

# 最初にメッセージを表示する
print("質問を入力してください")

conversation_history = []

while True:
    # ユーザーの入力した文字を変数「user_input」に格納
    user_input = input()

    # ユーザーの入力した文字が「exit」の場合はループを抜ける
    if user_input == "exit":
        break

    # ユーザーの質問を会話の履歴に追加
    conversation_history.append({"role": "user", "content": user_input})

    # 別の関数から回答を取得
    answer = answer_question(user_input, conversation_history)

    # ChatGPTの回答を会話履歴に追加
    conversation_history.append({"role": "assistant", "content": answer})

    print("ChatGPT:", answer)