dev/modify-embedding-test (#4)

- try to change embedding model (6125df0f09d11d466fb34cc188e156a8b8d2d7e0)
- Embeddings multilingual-e5-largeとLLM rinnaを使えるようにした (143e47c91fdbdedb218bb197b96d4a6ed84d892d)

app.py

@@ -1,33 +1,105 @@
+# from time import time
 import gradio as gr
 from langchain.chains import RetrievalQA
 from langchain.embeddings import OpenAIEmbeddings
-from langchain.llms import OpenAI
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain.llms import HuggingFacePipeline
+
+# from langchain.llms import OpenAI
 from langchain.chat_models import ChatOpenAI
 from langchain.vectorstores import Qdrant
 from openai.error import InvalidRequestError
 from qdrant_client import QdrantClient
-from config import DB_CONFIG
+from config import DB_CONFIG, DB_E5_CONFIG
 
 
-PERSIST_DIR_NAME = "nvdajp-book"
-# MODEL_NAME = "text-davinci-003"
-# MODEL_NAME = "gpt-3.5-turbo"
-# MODEL_NAME = "gpt-4"
+def _get_config_and_embeddings(collection_name: str | None) -> tuple:
+    if collection_name is None or collection_name == "E5":
+        db_config = DB_E5_CONFIG
+        model_name = "intfloat/multilingual-e5-large"
+        model_kwargs = {"device": "cpu"}
+        encode_kwargs = {"normalize_embeddings": False}
+        embeddings = HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs=model_kwargs,
+            encode_kwargs=encode_kwargs,
+        )
+    elif collection_name == "OpenAI":
+        db_config = DB_CONFIG
+        embeddings = OpenAIEmbeddings()
+    else:
+        raise ValueError("Unknow collection name")
+    return db_config, embeddings
 
 
-def get_retrieval_qa(model_name: str | None, temperature: int, option: str | None) -> RetrievalQA:
-    embeddings = OpenAIEmbeddings()
-    db_url, db_api_key, db_collection_name = DB_CONFIG
-    client = QdrantClient(url=db_url, api_key=db_api_key)
-    db = Qdrant(client=client, collection_name=db_collection_name, embeddings=embeddings)
+def _get_rinna_llm(temperature: float):
+    model = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
+    tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False)
+    model = AutoModelForCausalLM.from_pretrained(
+        model,
+        load_in_8bit=True,
+        torch_dtype=torch.float16,
+        device_map="auto",
+    )
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=1024,
+        temperature=temperature,
+    )
+    llm = HuggingFacePipeline(pipeline=pipe)
+    return llm
+
+
+def _get_llm_model(
+    model_name: str | None,
+    temperature: float,
+):
     if model_name is None:
-        model = "gpt-3.5-turbo"
+        model = "rinna"
+    elif model_name == "rinna":
+        model = "rinna"
     elif model_name == "GPT-3.5":
         model = "gpt-3.5-turbo"
     elif model_name == "GPT-4":
         model = "gpt-4"
     else:
-        model = "gpt-3.5-turbo"
+        raise ValueError("Unknow model name")
+    if model.startswith("gpt"):
+        llm = ChatOpenAI(model=model, temperature=temperature)
+    elif model == "rinna":
+        llm = _get_rinna_llm(temperature)
+    return llm
+
+
+# prompt_template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
+
+# {context}
+
+# Question: {question}
+# Answer in Japanese:"""
+# PROMPT = PromptTemplate(
+#     template=prompt_template, input_variables=["context", "question"]
+# )
+
+
+def get_retrieval_qa(
+    collection_name: str | None,
+    model_name: str | None,
+    temperature: float,
+    option: str | None,
+) -> RetrievalQA:
+    db_config, embeddings = _get_config_and_embeddings(collection_name)
+    db_url, db_api_key, db_collection_name = db_config
+    client = QdrantClient(url=db_url, api_key=db_api_key)
+    db = Qdrant(
+        client=client, collection_name=db_collection_name, embeddings=embeddings
+    )
+
     if option is None or option == "All":
         retriever = db.as_retriever()
     else:
@@ -36,15 +108,19 @@ def get_retrieval_qa(model_name: str | None, temperature: int, option: str | Non
                 "filter": {"category": option},
             }
         )
-    return RetrievalQA.from_chain_type(
-        llm=ChatOpenAI(
-            model=model,
-            temperature=temperature
-        ),
+
+    llm = _get_llm_model(model_name, temperature)
+
+    # chain_type_kwargs = {"prompt": PROMPT}
+
+    result = RetrievalQA.from_chain_type(
+        llm=llm,
         chain_type="stuff",
         retriever=retriever,
         return_source_documents=True,
+        # chain_type_kwargs=chain_type_kwargs,
     )
+    return result
 
 
 def get_related_url(metadata):
@@ -60,8 +136,10 @@ def get_related_url(metadata):
         yield f'<p>URL: <a href="{url}">{url}</a> (category: {category})</p>'
 
 
-def main(query: str, model_name: str, option: str, temperature: int):
-    qa = get_retrieval_qa(model_name, temperature, option)
+def main(
+    query: str, collection_name: str, model_name: str, option: str, temperature: float
+):
+    qa = get_retrieval_qa(collection_name, model_name, temperature, option)
     try:
         result = qa(query)
     except InvalidRequestError as e:
@@ -77,9 +155,14 @@ nvdajp_book_qa = gr.Interface(
     fn=main,
     inputs=[
         gr.Textbox(label="query"),
-        gr.Radio(["GPT-3.5", "GPT-4"], label="Model", info="選択なしで「3.5」を使用"),
-        gr.Radio(["All", "ja-book", "ja-nvda-user-guide", "en-nvda-user-guide"], label="絞り込み", info="ドキュメント制限する?"),
-        gr.Slider(0, 2)
+        gr.Radio(["E5", "OpenAI"], label="Embedding", info="選択なしで「E5」を使用"),
+        gr.Radio(["rinna", "GPT-3.5", "GPT-4"], label="Model", info="選択なしで「rinna」を使用"),
+        gr.Radio(
+            ["All", "ja-book", "ja-nvda-user-guide", "en-nvda-user-guide"],
+            label="絞り込み",
+            info="ドキュメント制限する?",
+        ),
+        gr.Slider(0, 2),
     ],
     outputs=[gr.Textbox(label="answer"), gr.outputs.HTML()],
 )

config.py

@@ -4,18 +4,24 @@ import os
 SAAS = True
 
 
-def get_db_config():
-    url = os.environ["QDRANT_URL"]
+def get_db_config(cname="nvdajp-book"):
+    if cname == "nvdajp-book":
+        url = os.environ["QDRANT_URL"]
+    elif cname == "nvdajp-book-e5":
+        url = os.environ["QDRANT_E5_URL"]
     api_key = os.environ["QDRANT_API_KEY"]
-    collection_name = "nvdajp-book"
+    collection_name = cname
     return url, api_key, collection_name
 
 
-def get_local_db_congin():
+def get_local_db_congin(cname="nvdajp-book"):
     url = "localhost"
     # api_key = os.environ["QDRANT_API_KEY"]
-    collection_name = "nvdajp-book"
+    collection_name = cname
     return url, None, collection_name
 
 
 DB_CONFIG = get_db_config() if SAAS else get_local_db_congin()
+DB_E5_CONFIG = (
+    get_db_config("nvdajp-book-e5") if SAAS else get_local_db_congin("nvdajp-book-e5")
+)

requirements.txt

@@ -4,3 +4,5 @@ tiktoken
 gradio
 qdrant-client
 beautifulsoup4
+accelerate
+bitsandbytes

store.py

@@ -1,10 +1,12 @@
 from langchain.document_loaders import ReadTheDocsLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.embeddings import OpenAIEmbeddings
+from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import Qdrant
+
 # from qdrant_client import QdrantClient
 from nvda_ug_loader import NVDAUserGuideLoader
-from config import DB_CONFIG
+from config import DB_CONFIG, DB_E5_CONFIG
 
 
 CHUNK_SIZE = 500
@@ -23,37 +25,55 @@ def get_documents(path: str):
     for doc in docs:
         org_metadata = doc.metadata
         source = _remove_prefix_path(org_metadata["source"])
-        add_meta = {"category": category, "source": source, "url": f"{base_url}{source}"}
+        add_meta = {
+            "category": category,
+            "source": source,
+            "url": f"{base_url}{source}",
+        }
         doc.metadata = org_metadata | add_meta
         yield doc
 
 
 def get_text_chunk(docs):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=0)
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=CHUNK_SIZE, chunk_overlap=0
+    )
     texts = text_splitter.split_documents(docs)
     return texts
 
 
-def store(texts):
-    embeddings = OpenAIEmbeddings()
-    db_url, db_api_key, db_collection_name = DB_CONFIG
-    # client = QdrantClient(url=db_url, api_key=db_api_key, prefer_grpc=True)
+def store(texts, mname):
+    if mname == "openai":
+        embeddings = OpenAIEmbeddings()
+        db_url, db_api_key, db_collection_name = DB_CONFIG
+    elif mname == "e5":
+        model_name = "intfloat/multilingual-e5-large"
+        model_kwargs = {"device": "cuda"}
+        encode_kwargs = {"normalize_embeddings": False}
+        embeddings = HuggingFaceEmbeddings(
+            model_name=model_name,
+            model_kwargs=model_kwargs,
+            encode_kwargs=encode_kwargs,
+        )
+        db_url, db_api_key, db_collection_name = DB_E5_CONFIG
+    else:
+        raise ValueError("Invalid mname")
     _ = Qdrant.from_documents(
         texts,
         embeddings,
         url=db_url,
         api_key=db_api_key,
-        collection_name=db_collection_name
+        collection_name=db_collection_name,
     )
 
 
-def rtd_main(path: str):
+def rtd_main(path: str, mname: str):
     docs = get_documents(path)
     texts = get_text_chunk(docs)
-    store(texts)
+    store(texts, mname)
 
 
-def nul_main(url: str):
+def nul_main(url: str, mname: str):
     if "www.nvda.jp" in url:
         category = "ja-nvda-user-guide"
     else:
@@ -61,25 +81,27 @@ def nul_main(url: str):
     loader = NVDAUserGuideLoader(url, category)
     docs = loader.load()
     texts = get_text_chunk(docs)
-    store(texts)
+    store(texts, mname)
 
 
 if __name__ == "__main__":
     """
-    $ python store.py rtd "data/rtdocs/nvdajp-book.readthedocs.io/ja/latest"
-    $ python store.py nul "https://www.nvaccess.org/files/nvda/documentation/userGuide.html"
-    $ python store.py nul "https://www.nvda.jp/nvda2023.1jp/ja/userGuide.html"
+    $ python store.py rtd "data/rtdocs/nvdajp-book.readthedocs.io/ja/latest" openai
+    $ python store.py nul "https://www.nvaccess.org/files/nvda/documentation/userGuide.html" e5
+    $ python store.py nul "https://www.nvda.jp/nvda2023.1jp/ja/userGuide.html" e5
     """
     import sys
+
     args = sys.argv
-    if len(args) != 3:
+    if len(args) != 4:
         print("No args, you need two args for type, html_path")
     else:
         type_ = args[1]
         path = args[2]
+        mname = args[3]
         if type_ == "rtd":
-            rtd_main(path)
+            rtd_main(path, mname)
         elif type_ == "nul":
-            nul_main(path)
+            nul_main(path, mname)
         else:
             print("No type for store")