Spaces:
Paused
Paused
File size: 5,372 Bytes
8034497 6ab28e5 2c70642 9022e07 6ab28e5 2c70642 6ab28e5 8034497 2c70642 8034497 2c70642 6ab28e5 2c70642 8034497 2c70642 9022e07 2c70642 9022e07 2c70642 99d3f35 2c70642 9022e07 2c70642 6ab28e5 2c70642 6ab28e5 9bc4a6c 6ab28e5 9bc4a6c 6ab28e5 9bc4a6c 6ab28e5 2c70642 8034497 2c70642 6ab28e5 8034497 6ab28e5 99d3f35 2c70642 99d3f35 6ab28e5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
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
E5_MODEL_NAME = "intfloat/multilingual-e5-large"
E5_MODEL_KWARGS = {"device": "cuda:0" if torch.cuda.is_available() else "cpu"}
E5_ENCODE_KWARGS = {"normalize_embeddings": False}
E5_EMBEDDINGS = HuggingFaceEmbeddings(
model_name=E5_MODEL_NAME,
model_kwargs=E5_MODEL_KWARGS,
encode_kwargs=E5_ENCODE_KWARGS,
)
RINNA_MODEL_NAME = "rinna/bilingual-gpt-neox-4b-instruction-ppo"
RINNA_TOKENIZER = AutoTokenizer.from_pretrained(RINNA_MODEL_NAME, use_fast=False)
RINNA_MODEL = AutoModelForCausalLM.from_pretrained(
RINNA_MODEL_NAME,
load_in_8bit=True,
torch_dtype=torch.float16,
device_map="auto",
)
def _get_config_and_embeddings(collection_name: str | None) -> tuple:
if collection_name is None or collection_name == "E5":
db_config = DB_E5_CONFIG
embeddings = E5_EMBEDDINGS
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):
pipe = pipeline(
"text-generation",
model=RINNA_MODEL,
tokenizer=RINNA_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
):
now = time()
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"]]
sec_html = f"<p>実行時間: {(time() - now):.2f}秒</p>"
html = "<div>" + sec_html + "\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()
|