nvdajp-book-qa / app.py
terapyon's picture
dev/modify-embedding-test (#4)
2c70642
raw
history blame
5.22 kB
# from time import time
import gradio as gr
from langchain.chains import RetrievalQA
from langchain.embeddings import OpenAIEmbeddings
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, DB_E5_CONFIG
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_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 = "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:
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:
retriever = db.as_retriever(
search_kwargs={
"filter": {"category": option},
}
)
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):
urls = set()
for m in metadata:
# p = m['source']
url = m["url"]
if url in urls:
continue
urls.add(url)
category = m["category"]
# print(m)
yield f'<p>URL: <a href="{url}">{url}</a> (category: {category})</p>'
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:
return "回答が見つかりませんでした。別な質問をしてみてください", str(e)
else:
metadata = [s.metadata for s in result["source_documents"]]
html = "<div>" + "\n".join(get_related_url(metadata)) + "</div>"
return result["result"], html
nvdajp_book_qa = gr.Interface(
fn=main,
inputs=[
gr.Textbox(label="query"),
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()],
)
nvdajp_book_qa.launch()