AIを使ってアプリケーション名からCPEを求められないか？ #2 RAG

# はじめに

[前回](https://qiita.com/iwanagat85/items/7f49bca17e1379caa465)はLoRAでファインチューニングしてCPEを求めてみました。
ただ精度はいまいちですし、新しいCPEに対応しません。
RAGを用いて外部情報を参照してCPEを求めるようにしてみようと思います。


# 作成物
https://github.com/iwanagat85/cpe-generator/tree/lora%2Brag

# 目標

- RAGで外部情報を参照させる
- 精度が高くなったらいいな
- 入力例:
  `Visual Studio Code　0.2.9`
- 出力例:
  `cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*`


# VectorStoreの作成

埋め込みモデルは多言語対応している [Multilingual-E5-base](https://huggingface.co/intfloat/multilingual-e5-base) を使ってみます。
VectorStoreの操作には扱いやすそうな [Chroma](https://docs.trychroma.com/docs/overview/getting-started) を使ってみます。
さっそく `langchain_community.document_loaders.CSVLoader` を使ってドキュメントを読み込んでみます。


```python
documents = CSVLoader(csv_path, encoding="utf-8").load()
embedding = HuggingFaceEmbeddings(
    model_name="intfloat/multilingual-e5-base",
)
Chroma.from_documents(
    documents=documents,
    embedding=embedding,
    persist_directory="vectorstore",
)
```

```bash
$ python create_vectorstore.py
```

# VectorStoreの確認
163万件のドキュメントを読み込ませたところ1時間ほどかかりました。
永続化したVectorStoreを読み込ませて類似度85％、上位3件を検索してみます。

```python
embeddings=HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-base")
vectorstore = Chroma(persist_directory="vectorstore", embedding_function=embeddings)
retriever = vectorstore.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "k": 3,
        "score_threshold": 0.85,
    },
)
retrieved_docs = retriever.invoke(query)
```

なかなか良さそうです。

```text
retrieved_docs: [
  Document(
    id='94dd712f-9c16-4779-9089-538fddc99966',
    metadata={'row': 1303638, 'source': 'datas/categorized_cpes.csv'},
    page_content='title: Microsoft Visual Studio Code 0.2.9\npart: a\nvendor: microsoft\nproduct: visual_studio_code\nversion: 0.2.9'
  ),
  Document(
    id='33300921-f7f9-41e0-abf8-080befa9b1b5',
    metadata={'row': 850060, 'source': 'datas/categorized_cpes.csv'},
    page_content='title: Microsoft Visual Studio Code 0.2.9 for Python\npart: a\nvendor: microsoft\nproduct: visual_studio_code\nversion: 0.2.9'
  ),
  Document(
    id='9975855f-14d6-4f73-89fe-5128d2befca5',
    metadata={'row': 689877, 'source': 'datas/categorized_cpes.csv'},
    page_content='title: Microsoft Visual Studio Code 0.20.0\npart: a\nvendor: microsoft\nproduct: visual_studio_code\nversion: 0.20.0'
  )
]
```


# インデックスをプロンプトテンプレートに挿入

VectoreStoreから得られた情報をプロンプトに挿入し、その情報を参考するよう指示します。
[LangChain](https://www.langchain.com/langchain)では[LCEL](https://python.langchain.com/docs/expression_language/)が推奨されていますが今回は自力で。

```diff_python
prompt_template = PromptTemplate(
    template=(
        "Generate a JSON from the given text.\n"
        "{format_instructions}\n\n"
+       "Please refer to the information below.\n\n"
+       "### Following information:\n"
+       "{context}\n\n"
    ),
-   input_variables=[],
+   input_variables=["context"],
    partial_variables={"format_instructions": format_instructions}
)
chat_template = tokenizer.apply_chat_template(
    [
-       {"role": "system", "content": prompt_template.format()},
+       {"role": "system", "content": prompt_template.format(context=retrieved_docs)},
        {"role": "user", "content": query}
    ],
    tokenize=False,
    add_generation_prompt=True
)
```


# RAGを使用してCPEを生成してみる

````bash
$ generate.py 
Input: Visual Studio Code 0.2.9
Output:
#0: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
#1: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
#2: ```json
"a"    "microsoft"  "visual_studio_code"  "0.2.9"
```
Got invalid return object. Expected key `part` to be present, but got a
For troubleshooting, visit: https://python.langchain.com/docs/troubleshooting/errors/OUTPUT_PARSING_FAILURE 
#3: ```json
"a"    "microsoft"    "visual_studio_code"    "0.2.9"
```
Got invalid return object. Expected key `part` to be present, but got a
For troubleshooting, visit: https://python.langchain.com/docs/troubleshooting/errors/OUTPUT_PARSING_FAILURE 
#4: ```json
"a"    "microsoft"    "visual_studio_code"    "0.2.9"
```
Got invalid return object. Expected key `part` to be present, but got a
For troubleshooting, visit: https://python.langchain.com/docs/troubleshooting/errors/OUTPUT_PARSING_FAILURE 
#5: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
#6: ```json
"a"      "microsoft" "visual_studio_code" "0.2.9"
```
Got invalid return object. Expected key `part` to be present, but got a
For troubleshooting, visit: https://python.langchain.com/docs/troubleshooting/errors/OUTPUT_PARSING_FAILURE 
#7: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
#8: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
#9: cpe:2.3:a:microsoft:visual_studio_code:0.2.9:*:*:*:*:*:*:*
````

パースに失敗するケースが増えてしまいましたが、他6件はすべて同じで正しいCPEを出力しています。
LoRAのみの時に出現していたベンダー名 `ms`、`ms-vim` などが出なくなりました。

# まとめ

LoRA+RAGで精度が上がった（ように見えます）。
`CSVLoader` に全て任せたVectorStoreのドキュメントなど改良の余地は色々とありそうです。
あと、パースに失敗するケースが増えたの何故でしょう？
まだまだ味がします。