localtest.py

from backend import Backend
import frontend
import numpy as np
from pandas import DataFrame
import pickle

with open("./himitsudogu_db.pkl", "rb") as file:
    himitsudogu_db: dict = pickle.load(file)

class LocaltestBackend(Backend):
    """
    ローカル開発環境にSentence BERTを連れ込んでくるのは無理なので、
    文章をベクトル化する行程は乱数発生とする。
    """

    def __init__(self):
        super().__init__()
        self.feature_matrix = himitsudogu_db["feature_matrix_s"][
            "sonoisa/sentence-bert-base-ja-mean-tokens-v2"
        ]

    def on_submit_button_press(self, query: str) -> DataFrame:
        query_feature_vector = np.random.random((768,))
        cs_s = self.feature_matrix @ query_feature_vector
        ranked_index_s = np.argsort(cs_s)[::-1]
        output = DataFrame(columns=["類似度", "名前", "説明"])
        for rank, i in enumerate(ranked_index_s[:20], 1):
            output.loc[rank] = [
                cs_s[i],
                himitsudogu_db["name_s"][i],
                himitsudogu_db["description_s"][i],
            ]
        return output

frontend.launch_frontend(backend=LocaltestBackend())