gpt-index(0.2.5)をOpenAI APIなし＆日本語で動かす

# gpt-index default_prompts.py
DEFAULT_TEXT_QA_PROMPT_TMPL = (
    "Context information is below. \n"
    "---------------------\n"
    "{context_str}"
    "\n---------------------\n"
    "Given the context information and not prior knowledge, "
    "answer the question: {query_str}\n"
)

DEFAULT_PROMPT = """
文脈情報は以下です。
---
{context_str}
---
事前知識ではなく、文脈情報を参考に質問に答えてください。：{query_str}
"""

# gpt-index default_prompts.py DEFAULT_REFINE_PROMPT_TMPL の日本語版
REFINE_PROMPT = """
質問は以下です。：{query_str}
すでに答えの候補があります。：{existing_answer}
必要な場合のみ、以下の文脈情報を使ってこの答えを改良することができます。
---
{context_msg}
---
この文脈情報により、元の答えを改良して質問に答えてください。
文脈情報が有用でない場合は元の答えをそのまま返してください。
"""

from transformers import T5Tokenizer, AutoModelForCausalLM

tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt-1b")
qa_model = AutoModelForCausalLM.from_pretrained("./models/qa-refine-japanese-gpt-1b")

def generate(prompt):

    token_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
    n = len(token_ids[0])

    with torch.no_grad():
        output_ids = qa_model.generate(
            token_ids.to(qa_model.device),
            pad_token_id=tokenizer.pad_token_id,
            bos_token_id=tokenizer.bos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    # 出力から回答部分のみを切り出す
    output = tokenizer.decode(output_ids.tolist()[0][n:])
    return output.replace("</s>", "")

from gpt_index import LLMPredictor
from langchain.llms.base import LLM
from typing import Optional, List, Mapping, Any
class CustomLLM(LLM):
        
    @property
    def _llm_type(self) -> str:
        return "custom"
    
    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        return generate(prompt)
    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        return {"name":"custom"}
llm = CustomLLM()
llm_predictor = LLMPredictor(llm)

from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from gpt_index import LangchainEmbedding
from sentence_transformers import SentenceTransformer

# いったんall-mpnet-base-v2で初期化
embed_model = LangchainEmbedding(HuggingFaceEmbeddings(model_name="oshizo/sbert-jsnli-luke-japanese-base-lite"))
# max_seq_length上書き
embed_model._langchain_embedding.client.max_seq_length = 256

import urllib
import json
with urllib.request.urlopen("https://ja.wikipedia.org/wiki/%E3%81%BC%E3%81%A3%E3%81%A1%E3%83%BB%E3%81%96%E3%83%BB%E3%82%8D%E3%81%A3%E3%81%8F!?action=cirrusdump") as f:
    data = f.read()
text = json.loads(data)[0]["_source"]["text"]

from gpt_index import GPTSimpleVectorIndex
from gpt_index.readers.schema.base import Document

documents = []
for i in range(0, len(text), 200):
    documents.append(Document(text[i:i+200]))
    if i != 0:
        documents.append(Document(text[i-100:i+100]))

from gpt_index import GPTSimpleVectorIndex

index = GPTSimpleVectorIndex(documents, llm_predictor=llm_predictor, embed_model=embed_model)

from gpt_index.prompts.prompts import RefinePrompt, QuestionAnswerPrompt
refine_prompt = RefinePrompt(REFINE_PROMPT)
default_prompt = QuestionAnswerPrompt(DEFAULT_PROMPT)

response = index.query(
    "虹夏ちゃんのお姉さんの職業はなに？",  
    mode="embedding", 
    verbose=True, 
    embed_model=embed_model,
    text_qa_template=default_prompt,
    refine_template=refine_prompt,
    similarity_top_k=3
)
print(response)

> Top 3 nodes:
> [Node 40aa0d84-e8e6-4102-8bab-ce864fc34bf2] て作曲の経験もあり、ひとりにアドバイスを与えた。 ひとりの動画投稿サイトでの収益を貯金していて、ギターを壊してしまったひとりにそのお金を与えている。 目はいつも何か（髪の毛や、ふたりの手など）に...
> [Node ea6fdcea-cd80-4696-af1f-d9ca8334aa00] 束バンドの後藤ひとり」であることは公にはされていない。 伊地知 虹夏（いじち にじか） 声 - 鈴代紗弓 誕生日：5月29日 / 血液型：A型 ドラム担当。下北沢高校2年→3年。明るく世話焼きで...
> [Node 7a1d891a-4450-4269-b23c-c7077ee9feed] に遠ざけられ、結局ひとりが作曲したものをリョウが編曲することになった。 虹夏と喜多によると（普段の奇行から忘れがちだが）美少女。文化祭のクラスの出し物でメイド服を着させられたときは、皆に似合って...
> Searching in chunk: て作曲の経験もあり、ひとりにアドバイスを与えた。 ひとりの動画投稿サイトでの収益を貯金していて...
> Searching in chunk: 束バンドの後藤ひとり」であることは公にはされていない。 伊地知 虹夏（いじち にじか） 声 -...
> Searching in chunk: に遠ざけられ、結局ひとりが作曲したものをリョウが編曲することになった。 虹夏と喜多によると（普...
> Initial response: 後藤 美智代
> Refine context: 束バンドの後藤ひとり」であることは公にはされていない。 伊地知 虹夏（いじち にじか） 声 -...
> Refined response: ライブハウスの店長
> Refine context: に遠ざけられ、結局ひとりが作曲したものをリョウが編曲することになった。 虹夏と喜多によると（普...
> Refined response: ライブハウスの店長
> [query] Total LLM token usage: 1436 tokens
> [query] Total embedding token usage: 30 tokens
'ライブハウスの店長'

response = index.query(
    "この作品のあらすじを教えて？", 
    mode="embedding", 
    verbose=True, 
    embed_model=embed_model,
    text_qa_template=default_prompt,
    refine_template=refine_prompt,
    similarity_top_k=3
)
response.response

> Top 1 nodes:
> [Node bf228da7-5efb-4322-a50f-e220f38b680b] 込む際に、ひとりの会話が視聴者に聞かせるための内容か音として鳴っているだけのものか、精査しながら場面構成を整えていく手法がとられている。また、原作に沿った内容でありながらもギャグの勢いをアニメだ...
> Searching in chunk: 込む際に、ひとりの会話が視聴者に聞かせるための内容か音として鳴っているだけのものか、精査しなが...
> Initial response: ギャグの勢いをアニメだからこそできる表現
> [query] Total LLM token usage: 411 tokens
> [query] Total embedding token usage: 21 tokens
'ギャグの勢いをアニメだからこそできる表現'

response = index.query(
    "後藤ひとりがギターに熱中するようになった理由は？", 
    mode="embedding", 
    verbose=True, 
    embed_model=embed_model,
    text_qa_template=default_prompt,
    refine_template=refine_prompt,
    similarity_top_k=3
)
response.response
> Top 3 nodes:
> [Node 43ee6b50-eb6b-444f-ad8c-621809d48f6a] の時に、暗い学生時代から一転してスターとなったバンドマンのインタビューを目にしたことで、父親から借りたギターに没頭する。毎日6時間以上の練習を約3年間欠かさず行ってきたため、個人としての演奏技術...
> [Node 5f1d20a4-5dfa-4335-83cd-e4c5063be3e4] 「心が乙女なおっさん」。好きな音楽ジャンルはパンクロック。 ぽいずん♡やみ 誕生日：2月14日 / 血液型：B型 音楽情報サイトに寄稿するフリーライター。本名は佐藤 愛子（さとう あいこ）。23...
> [Node 39f627ae-846e-418b-9c9b-348be5e1c7ac] ィング音源を使用しているほか、ひとりが行った奏法で演奏している。この楽曲は演奏中にひとりの1弦が切れてしまうシーンで使用されているため、この展開を踏まえて編曲を行い、ギターソロ後は1弦を使わなく...
> Searching in chunk: の時に、暗い学生時代から一転してスターとなったバンドマンのインタビューを目にしたことで、父親か...
> Searching in chunk: 「心が乙女なおっさん」。好きな音楽ジャンルはパンクロック。 ぽいずん♡やみ 誕生日：2月14日...
> Searching in chunk: ィング音源を使用しているほか、ひとりが行った奏法で演奏している。この楽曲は演奏中にひとりの1弦...
> Initial response: 暗い学生時代から一転してスターとなったバンドマンのインタビューを目にしたこと
> Refine context: 「心が乙女なおっさん」。好きな音楽ジャンルはパンクロック。 ぽいずん♡やみ 誕生日：2月14日...
> Refined response: 癖
> Refine context: ィング音源を使用しているほか、ひとりが行った奏法で演奏している。この楽曲は演奏中にひとりの1弦...
> Refined response: 癖
> [query] Total LLM token usage: 1463 tokens
> [query] Total embedding token usage: 35 tokens
'癖'