chore: Add utils, config, embeddings, retrievers, prompt, and llm modules

app.py

@@ -1,63 +1,52 @@
+# app.py
+import os
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
-
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough, RunnableLambda
+from langchain_community.document_transformers import LongContextReorder
+from config import LLM_MODEL, STREAMING
+from embeddings import get_embeddings
+from retrievers import load_retrievers
+from llm import get_llm
+from prompt import get_prompt
+
+
+def create_rag_chain(chat_history):
+    embeddings = get_embeddings()
+    retriever = load_retrievers(embeddings)
+    llm = get_llm(streaming=STREAMING)
+    prompt = get_prompt(chat_history)
+
+    return (
+        {
+            "context": retriever
+            | RunnableLambda(LongContextReorder().transform_documents),
+            "question": RunnablePassthrough(),
+        }
+        | prompt
+        | llm.with_config(configurable={"llm": LLM_MODEL})
+        | StrOutputParser()
+    )
+
+
+def respond_stream(message, history):
+    rag_chain = create_rag_chain(history)
     response = ""
+    for chunk in rag_chain.stream(message):
+        response += chunk
+        yield response
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
 
-        response += token
-        yield response
+def respond(message, history):
+    rag_chain = create_rag_chain(history)
+    return rag_chain.invoke(message)
+
 
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
 demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+    respond_stream if STREAMING else respond,
+    title="판례에 대해서 물어보세요!",
+    description="안녕하세요!\n저는 판례에 대한 인공지능 QA봇입니다. 판례에 대해 깊은 지식을 가지고 있어요. 판례에 관한 도움이 필요하시면 언제든지 질문해주세요!",
 )
 
-
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()

config.py

@@ -0,0 +1,13 @@
+# config.py
+import os
+from dotenv import load_dotenv
+
+load_dotenv()
+
+FAISS_DB_INDEX = "./index_faiss"
+BM25_INDEX = "./index_bm25/kiwi.pkl"
+CHUNK_SIZE = 2000
+CHUNK_OVERLAP = 200
+EMBEDDING_MODEL = "BAAI/bge-m3"
+LLM_MODEL = os.getenv("MODEL_KEY", "gemini")
+STREAMING = os.getenv("STREAMING", "true").lower() == "true"

embeddings.py

@@ -0,0 +1,12 @@
+# embeddings.py
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from config import EMBEDDING_MODEL
+from utils import get_device
+
+
+def get_embeddings():
+    return HuggingFaceBgeEmbeddings(
+        model_name=EMBEDDING_MODEL,
+        model_kwargs={"device": get_device()},
+        encode_kwargs={"normalize_embeddings": True},
+    )

llm.py

@@ -0,0 +1,54 @@
+# llm.py
+from langchain_openai import ChatOpenAI
+from langchain_anthropic import ChatAnthropic
+from langchain_google_genai import GoogleGenerativeAI
+from langchain_groq import ChatGroq
+from langchain_community.chat_models import ChatOllama
+from langchain_core.runnables import ConfigurableField
+from langchain.callbacks.base import BaseCallbackHandler
+
+
+class StreamCallback(BaseCallbackHandler):
+    def on_llm_new_token(self, token: str, **kwargs):
+        print(token, end="", flush=True)
+
+
+def get_llm(streaming=True):
+    return ChatOpenAI(
+        model="gpt-4",
+        temperature=0,
+        streaming=streaming,
+        callbacks=[StreamCallback()],
+    ).configurable_alternatives(
+        ConfigurableField(id="llm"),
+        default_key="gpt4",
+        claude=ChatAnthropic(
+            model="claude-3-opus-20240229",
+            temperature=0,
+            streaming=streaming,
+            callbacks=[StreamCallback()],
+        ),
+        gpt3=ChatOpenAI(
+            model="gpt-3.5-turbo",
+            temperature=0,
+            streaming=streaming,
+            callbacks=[StreamCallback()],
+        ),
+        gemini=GoogleGenerativeAI(
+            model="gemini-1.5-flash",
+            temperature=0,
+            streaming=streaming,
+            callbacks=[StreamCallback()],
+        ),
+        llama3=ChatGroq(
+            model_name="llama3-70b-8192",
+            temperature=0,
+            streaming=streaming,
+            callbacks=[StreamCallback()],
+        ),
+        ollama=ChatOllama(
+            model="EEVE-Korean-10.8B:long",
+            streaming=streaming,
+            callbacks=[StreamCallback()],
+        ),
+    )

prompt.py

@@ -0,0 +1,31 @@
+# prompt.py
+from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
+
+PROMPT_TEMPLATE = """당신은 판사이자 20년 차 법률 전문가입니다. 주어진 질문에 대해 문서의 정보를 최대한 활용하여 답변하세요. 질문자는 자신의 상황을 설명할 것이며, 질문자의 상황과 비슷한 판례를 설명해줘야 합니다. 가장 최근 사건 순으로 소개하며, 초등학생이 이해할 수 있도록 최대한 자세하고 쉽게 설명하세요. 답변은 [사건명 1]..., [사건명 2]... 순서로 구성합니다. 문서에서 답변을 찾을 수 없는 경우, "문서에 답변이 없습니다."라고 답변하세요.
+
+답변의 출처(source)를 반드시 표기합니다. 출처는 메타데이터의 판례일련번호, 사건명, 사건번호 순으로 표기합니다.
+
+---
+
+# 주어진 문서:
+{context}
+
+# 질문: {question}
+
+# 답변:
+
+# 출처:
+- source1
+- source2
+- ...
+"""
+
+
+def get_prompt(chat_history):
+    return ChatPromptTemplate.from_messages(
+        [
+            ("system", PROMPT_TEMPLATE),
+            MessagesPlaceholder(variable_name="history"),
+            ("human", "{question}"),
+        ]
+    ).partial(history=chat_history.messages)

retrievers.py

@@ -0,0 +1,22 @@
+# retrievers.py
+import pickle
+from langchain.vectorstores import FAISS
+from langchain.retrievers import EnsembleRetriever
+from kiwipiepy import Kiwi
+from config import FAISS_DB_INDEX, BM25_INDEX
+
+
+def load_retrievers(embeddings):
+    faiss_db = FAISS.load_local(
+        FAISS_DB_INDEX, embeddings, allow_dangerous_deserialization=True
+    )
+    faiss_retriever = faiss_db.as_retriever(search_type="mmr", search_kwargs={"k": 10})
+
+    with open(BM25_INDEX, "rb") as f:
+        bm25_retriever = pickle.load(f)
+
+    return EnsembleRetriever(
+        retrievers=[bm25_retriever, faiss_retriever],
+        weights=[0.7, 0.3],
+        search_type="mmr",
+    )

utils.py

@@ -0,0 +1,11 @@
+# utils.py
+import torch
+
+
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda:0"
+    elif torch.backends.mps.is_available():
+        return "mps"
+    else:
+        return "cpu"