跳至主要内容

讓你的文件活起來 - RAG 實作

高見龍
五倍學院 負責人

這是在 iThome 所舉辦的 Hello World Dev Conference 工作坊的內容。本次工作坊將帶著大家理解 RAG(Retrieval-Augmented Generation)技術,並使用 Python 程式語言及 LangChain 框架進行實作,最終將文件轉化為可以進行互動問題的知識庫。

現在沒用過 AI 服務的人大概都要被歸類到上個世代的化石了。不得不說,這些 AI 服務真的很厲害,但到目前一直有個比較大的困擾,就是它有時候會一本正經的講幹話,就算遇到它不會的問題,因為它講出來的內容太有自信,導致不知道它到底是真的假的。

舉個例子,如果你幫公司做了一個用來做知識管理(Knowledge Management, KM)的網站,公司內部一些相關的規定都可以在這裡查到。你在這個網站上掛了一個聊天機器人,當訪客問機器人你們家的產品多少價錢或是該怎麼退換貨的時候,這時候不知道就該說不知道,而不是硬擠答案出來。我們可以透過適當的「提示(Prompt)」來限制 AI 的回答,但還是希望 AI 不要隨便亂回答,但就是因為這樣,才有了 RAG 的出現。

重要!請先看這裡!

因為活動現場網路頻寬可能不夠快,而且因為會用到的模型檔案有點大,如果大家報名這個工作坊而且想在現場跟著一起實作的話,建議可以找空檔完成以下步驟:

  1. 下載 Ollama 應用程式

請在這裡選擇合適的版本下載並安裝。

  1. 下載模型

在本次工作坊中,我將使用 mistral 模型做為範例,這個模型大概有 4GB,你也可以選擇其它模型。先執行指令下載 mistral 模型:

$ ollama pull mistral

另外還會用到另一個比較小的模型 nomic-embed-text,也可以順便先拉下來:

$ ollama pull nomic-embed-text

或是直接執行也可以:

$ ollama run mistral

如果還沒下載過 mistral 模型的話,這個指令會幫你下載,並且直接進入聊天模式,然後就可以開始問它問題了:

$ ollama run mistral
>>> 你好 :)
嗨,nice to meet you! 今天好吗?

(Hello, nice to meet you! How are you today?)

如果有需要帮助的话,我会尽力帮到您。

(If you need any help, I will do my best to assist.)

如果有空檔的話,以下這兩個也可以順便安裝一下:

  • Python:到 Python 官網下載合適版本,這個難度不高
  • 文字編輯器:本次工作坊會使用 VSCode,但你可使用自己順手的開發工具即可。

比較花時間的安裝到這裡差不多就算完成了。

名詞解釋

AI 時代一堆專有名詞或縮寫真的滿天飛,所以我們先從名詞解釋開始

LLM

LLM(Large Language Model),中文翻譯成「大型語言模型」,但這翻譯有跟沒有差不多。想想看,以前如果我們想要判斷使用者輸入了什麼內容,或是判斷語意,特別是中文,我們可能需要使用一些斷字、斷詞的資料庫才有辦法處理,特別是現在很多人「再在」「應因」不分的情況下,想要要判斷文字的意圖或情緒基本上是不太容易做到的。

雖然我們大部份的人,包括我自己也是,都不知道這到底是怎麼做到的,但 LLM 做到了。透過 LLM,可以幫我們解讀使用者的「輸入」,理解想要的問題,並且在適當的「提示(Prompt)」下回答使用者的問題,甚至有一些錯別字也不會影響結果。簡單來說,LLM 幫我們搞定輸入以及輸出的事情。

問題是,這個 LLM 的 M 是什麼 Model?又是怎麼訓練來的?像我們這種凡人應該是沒能力訓練模型的,大部份只能看那些神仙公司打架,然後撿他們打架掉出來的東西來用,像是 Facebook 的 Llama3、Google 的 Gemma、Mistral AI 的 mistral...等開源模型。

只是,就算我們撿到了這些訓練好的模型,大部份的我們大概也不太容易再對它再進行訓練,讓它更聰明。雖然我們可以透過更精準的 Prompt 請 LLM 回答的準確一點,但 LLM 並沒有我們想像中的聰明,它其實是很健忘的。舉個例子,當你你在跟 ChatGPT 聊天的過程中,感覺好像都會記得你跟它聊了什麼,這是因為 ChatGPT 在每次的對話裡,都會把同一個討論串的「情境」或「上下文(Context)」做個摘要,然後在下一個對話的時候一併傳給 ChatGPT。然而這樣不斷累積情境或上下文的做法是有限的,這東西又稱「上下文視窗(Context Window)」,這個窗戶沒辦法一直無限開下去,特別如果遇到你想要查詢的情境是一本或數本 PDF 電子書的時候,這可能就沒辦法了。

其次,就算是這些大公司訓練好的 Model,通常也都是根據現有的資料訓練的,如果想要問它最新一期大樂透開幾號,或是問它今天台積電的收盤價多少,它應該不會知道。另外,因為這些模型都是使用公開的資料進行訓練的,所以如果你想問它你公司的請假規定,照理它應該也不會知道。

所以在這種情況下,LLM 會怎麼回答你?還能怎麼回答,就瞎掰啊!

RAG

RAG(Retrieval-Augmented Generation),中文翻譯成「檢索增強生成」,RAG 的重點不在檢索(Retrieval)或生成(Generation),因為這本來 LLM 就能做到了,重在在於「增強(Argumented)」,透過它,LLM 可以更精準的回答我們的問題。

講這到這裡,可能會跟另一個名詞「微調(Fine-tuning)」有點像,但其實這兩個是不太一樣的概念。微調是指針對現有的 Model 再拿特定領域的資料進行訓練、調整,例如可餵食特定領域醫學相關的資料,增進這個 Model 在這方面的醫學知識。

而 RAG 並不是對原本的 Model 進行調整,而是在我們對 LLM 下 Prompt 的時候提供額外的參數或資訊,讓 LLM 能更準確的回答問題。簡單的說,RAG 是提供給 LLM「額外的知識」,再講的更白話一點,就是讓你原本對 LLM 下的 Prompt 更精準。

很多時候你以為可以透過「微調」來讓 LLM 聰明一點,事實上你可能需要的是 RAG。

Embedding

Embedding,中文翻譯成「嵌入」,這個詞可能比較難從字樣上來想像是怎麼回事。簡單來說,Embedding 是一種將資料轉換成「向量(Vector)」的過程或技術,這麼做可以讓 AI 更好地「理解」並比對資料的「相似性」。

什麼是「向量」?舉例來說,我想針對香蕉、芭樂、蘋果跟榴槤這四個水果的味道跟質地軟硬做個比較,如果我把味道跟質地用一個二維空間來呈現,可能看起來會像這樣:

臭 ^
  | 榴槤(0.9, 0.9)
  |
  |
  |
  |
  |
  |
  | 蘋果(0.5, 0.2) 芭樂(0.6, 0.2)
香 | 香蕉(0.3, 0.1)
+----------------------------------> x
軟 硬

透過幫這些水果「打分數」,就能幫它們在這個二維空間定出位置來。在二維空間越接近的水果,代表它們的味道跟質地越相似。

回到我們的主題,Embedding 技術可以讓我們把資料轉換成向量,這樣我們就可以透過向量的方式來比對資料的相似性。這裡指的向量不會像我們這裡只是二維的空間,而是可能有幾百維,甚至上千維。透過向量的方式來比對資料的相似性,可以除了能幫原本的 LLM 補充一些它原本不知道的事之夕,也讓它更容易找到精準的資料。

向量資料庫

如果每次都要對資料轉成向量,效能可能會有點差,所以有一個可以存放這些向量的地方是很重要的。也許你會想把這些向量另外存在文字檔裡,但這樣的話,每次要找相似的資料時,就要把所有的向量都讀進來,然後再逐一比對,這可能也會有效能上的問題。

還好現在有專門用來存放向量的資料庫,以我們這個工作坊會用到的 Chroma 為例,就是一個專門用來存放向量的資料庫,而且還有一些神奇的特異功能,讓我們不用懂太多演算法就能找出相似的向量。

動手作!

好啦,故事講的差不多了,是該時候動手做點東西了!

準備工作

程式語言及開發工具

雖然沒有限定程式語言,但以現在的大環境來看,Python 可能會是最好的選擇,並不是因為 Python 這個程式語言特別厲害,而是 Python 的生態系裡面有很多相關的套件可以使用,可以幫我們省下不少時間。

開發工具的話,因為我們要寫的程式碼也沒多複雜,所以用一般的文字編輯器就很夠了。如果你沒有慣用的文字編輯器,可以考慮使用 Visual Studio Code(簡稱 VSCode)。

環境安裝

可以從 Python 官網下載安裝包,或是透過 pyenv 來安裝也可以,細節可參閱環境安裝章節的介紹。

手工打造

在開始之前,我想先問試著問 Ollama 一個問題,看看它會怎麼回答:

$ ollama run mistral
>>> 請問高見龍是誰?
高見龍(高氏明龍,1592年-1645年)是中國明朝時期的一位知名詩人。他出生於江南,從小受教于家里,後入京立學。高見龍在文學界著名之作有《山水遺集
》、《詩集新選》等。他是明朝最後一位有名的北宋式詩人,對中國詩歌傳統產生了深刻影響。

喔喔喔,這麼厲害啊!我怎麼都不知道原來我是詩人?沒關係,待會我會借用這個頁面的資料並且透過 RAG 來改善這個結果。

為了保持每個專案的套件不互相干擾,我先透過 Python 內建 venv 建立並切換至虛擬環境:

$ python -m venv .venv
$ source .venv/bin/activate

對 Python 虛擬環境不熟悉的話,可參閱環境安裝的「使用 venv」小節介紹。再來,安裝待會要用的套件 ollama 以及用來計算向量相似度的 numpy

$ pip install ollama numpy

接著,我把「關於我」的頁面上的文字存成一個文字檔 about_me.txt 放在同一個目錄裡,接著建立一個 rag_v1.py 檔案,然後來寫點程式碼:

檔案:rag_v1.py
def parse_paragraph(filename):
with open(filename) as f:
return [line.strip() for line in f if line.strip()]


if __name__ == "__main__":
paragraphs = parse_paragraph("about_me.txt")

這個 parse_paragraph() 的功能其實滿簡單的,就是把 about_me.txt 檔案讀進來,去除多餘的空白行,並且讓每一行組裝成一個串列回傳回來。接著再借用 Ollama 套件幫我們把這個串列的每個元素轉換成向量:

檔案:rag_v1.py
import ollama

# ... 略 ...

def calc_embedings(paragraphs):
return [
ollama.embeddings(model="mistral", prompt=data)["embedding"]
for data in paragraphs
]

if __name__ == "__main__":
paragraphs = parse_paragraph("about_me.txt")
embeddings = calc_embedings(paragraphs)

把剛才處理好的段落傳給 calc_embedings() 函數,最終會得到一個計算好的向量串列。這裡我用了 mistral 模型來進行計算,如果要換成其它模型也可以。執行之後應該會發現,雖然 about_me.txt 的文字並不多,但跑起來會有明顯的卡頓感,以我 2021 年的 M1 Mac 筆記大概都要跑二、三秒才會跑完,這是因為計算這些向量就是需要時間(或需要更厲害的硬體),而且這個卡頓感可能會隨著餵食的資料越多會越明顯。這裡不能總是每次都這樣重算,所以接下來我會寫一個函數,在第一次計算完之後會先把計算好的向量存下來,下次如果還是算同一個檔案就不用再重算:

檔案:rag_v1.py
import ollama
import os
import json

# ... 略 ...

def cache_embeddings(filename, paragraphs):
embedding_file = f"cache/{filename}.json"

if os.path.isfile(embedding_file):
with open(embedding_file) as f:
return json.load(f)

os.makedirs("cache", exist_ok=True)

embeddings = calc_embedings(paragraphs)

with open(embedding_file, "w") as f:
json.dump(embeddings, f)

return embeddings

if __name__ == "__main__":
paragraphs = parse_paragraph("about_me.txt")
embeddings = cache_embeddings("about_me.txt", paragraphs)

這個 cache_embeddings() 函數大概就是把計算好的向量存起來,如果檔案已經存在就把它讀出來。但這個函數其實有點問題,就是如果 about_me.txt 的內容有調整的話,除非我把 cache/ 目錄裡的資料刪掉,不然不會重新計算。不過沒關係,這一段的手工打造的目的只是展示原理,晚點你會看到我用其它套件在處理的時候就不用煩惱這種瑣事,所以暫時先接受這個不完美。

第一次執行應該還是一樣會卡頓,但第二次執行之後,應該很明顯一下子就跑完了。

最後再補一個用來計算向量相似度的函數 calc_similar_vectors()

檔案:rag_v1.py
from numpy import linalg, dot

# ... 略 ...

def calc_similar_vectors(v, vectors):
v_norm = linalg.norm(v)
scores = [dot(v, item) / (v_norm * linalg.norm(item)) for item in vectors]
return sorted(enumerate(scores), reverse=True, key=lambda x: x[1])

# ... 略 ...

這個函數裡有些數學計算可能稍微複雜一點,基本上就是用來計算在向量們的距離。如前面提到的,向量越相似,表示這些向量的資料也越相似。這個函數回傳一個由依照相似度分數排序的索引值組成的串列,這樣我們就可以知道哪個向量最相似。

到這裡,前置工作就差不多準備好了,再來就是把所有的東西組合起來:

檔案:rag_v1.py

if __name__ == "__main__":
doc = "about_me.txt"
paragraphs = parse_paragraph(doc)
embeddings = cache_embeddings(doc, paragraphs)

prompt = input("請問你想問什麼問題?\n>>> ")

while prompt.lower() != "bye":
prompt_embedding = ollama.embeddings(model="mistral", prompt=prompt)[
"embedding"
]
similar_vectors = calc_similar_vectors(prompt_embedding, embeddings)[:3]

system_prompt = (
"現在開始使用我提供的情境來回答,只能使用繁體中文,不要有簡體中文字。如果你不確定答案,就說不知道。情境如下:"
+ "\n".join(paragraphs[vector[0]] for vector in similar_vectors)
)

response = ollama.chat(
model="mistral",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt},
],
)

print(response["message"]["content"])
prompt = input(">>> ")

來執行看看:

$ python rag_v1.py
您想問什麼問題?
>>> 高見龍是誰?
高見龍是一名網站開發者、講師、技術書作者、企業內訓及技術顧問。他是五倍學院(五倍紅寶石程式資訊教育股份有限公司)的負責人,並曾為 WebConf Taiwan 研討會發起人和主辦人,同時也是 PHPConf Taiwan 研討會發起人和主辦人。
>>> bye

Good,這看起來比剛才的詩人好多了!

完整程式碼:

檔案:rag_v1.py
import ollama
import os
import json
from numpy import linalg, dot

def parse_paragraph(filename):
with open(filename) as f:
return [line.strip() for line in f if line.strip()]

def calc_embedings(paragraphs):
return [
ollama.embeddings(model="mistral", prompt=data)["embedding"]
for data in paragraphs
]

def cache_embeddings(filename, paragraphs):
embedding_file = f"cache/{filename}.json"

if os.path.isfile(embedding_file):
with open(embedding_file) as f:
return json.load(f)

os.makedirs("cache", exist_ok=True)

embeddings = calc_embedings(paragraphs)

with open(embedding_file, "w") as f:
json.dump(embeddings, f)

return embeddings

def calc_similar_vectors(v, vectors):
v_norm = linalg.norm(v)
scores = [dot(v, item) / (v_norm * linalg.norm(item)) for item in vectors]
return sorted(enumerate(scores), reverse=True, key=lambda x: x[1])

if __name__ == "__main__":
doc = "about_me.txt"
paragraphs = parse_paragraph(doc)
embeddings = cache_embeddings(doc, paragraphs)

prompt = input("您想問什麼問題?\n>>> ")

while prompt.lower() != "bye":
prompt_embedding = ollama.embeddings(model="mistral", prompt=prompt)[
"embedding"
]
similar_vectors = calc_similar_vectors(prompt_embedding, embeddings)[:3]

system_prompt = (
"現在開始使用我提供的情境來回答,只能使用繁體中文,不要有簡體中文字。如果你不確定答案,就說不知道。情境如下:"
+ "\n".join(paragraphs[vector[0]] for vector in similar_vectors)
)

response = ollama.chat(
model="mistral",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": prompt},
],
)

print(response["message"]["content"])
prompt = input(">>> ")

使用 LangChain

上面這樣的手工打造的目的主要是理解 RAG 的原理,但實際做的時候不用這麼辛苦啦,Python 世界有很多厲害的善心人士幫我們做好了現成的套件,踩在這些巨人的肩膀上工作比較有效率。接下來就來看看怎麼使用 LangChain 幫我們更快搞定這些事。

不想建立新的虛擬環境的話,可在同一個目錄底下執行安裝套件:

$ pip install langchain langchain_community langchain_chroma pypdf

langchain 開頭的套件大概能猜出是為什麼要安裝,而最後的 pyddf 套件是晚點用到它解讀 PDF 用,可以順便一起裝一下。建立一個新的檔案 rag_v2.py,我們從頭開始寫:

檔案:rag_v2.py
from langchain_community.llms import Ollama
from langchain_community.document_loaders import TextLoader

llm = Ollama(model="mistral")

loader = TextLoader("about_me.txt")

這裡使用了 Langchain 社群提供的第三方套件來建立 LLM 物件,同時例用 TextLoader 來讀取 about_me.txt 的內容,就不用像我們前面自己手動呼叫 open() 函數讀取檔案。接著,我們來把這些文字分割成小段落:

檔案:rag_v2.py
from langchain.text_splitter import RecursiveCharacterTextSplitter

# ... 略 ...

text_splitter = RecursiveCharacterTextSplitter(
chunk_size=20,
chunk_overlap=5,
separators=[" ", ",", "\n"],
)

splited_docs = text_splitter.split_documents(loader.load())

這裡借用了 RecursiveCharacterTextSplitter 來把文字分割成小段落,這個做法跟前面我們自己手寫的 parse_paragraph() 函數差不多概念,但功能強多了。再來,把這些段落轉換成向量:

檔案:rag_v2.py
from langchain_community.embeddings import OllamaEmbeddings

# ... 略 ...

embeddings = OllamaEmbeddings(model="nomic-embed-text")

這個步驟跟前面自己寫的 calc_embedings() 函數差不多,只是這裡我們用了 OllamaEmbeddings 來幫我們計算向量。這裡我改用 nomic-embed-text 模型是因為用這個 Model 來計算速度比較快,當然你想繼續使用 mistral 也沒是可以。如果前面沒有安裝 nomic-embed-text 模型的話,這裡會需要請 Ollama 把模型拉下來,還好這個模型比 mistral 小多了,大概只有二百多 MB:

$ ollama pull nomic-embed-text

回來程式碼,來把這些算好的向量存起來:

檔案:rag_v2.py
from langchain_community.vectorstores import Chroma

# ... 略 ...

vector_db = Chroma.from_documents(
documents=splited_docs,
embedding=embeddings,
persist_directory="db",
collection_name="about",
)

這個步驟跟我們手刻的 cache_embeddings() 函數差不多,只是這裡我們用了 Chroma 來幫我們存放這些向量。persist_directory 是用來指定存放向量的目錄,我暫時把它設定成 db,你想改成什麼都可以,待會程式執行的時候就會看到在專案目錄底下自動生出這個目錄。而 collection_name 則是用來指定這些向量的集合名稱。

最後,把這些向量變成檢索器:

檔案:rag_v2.py
from langchain_core.prompts import ChatPromptTemplate

# ... 略 ...

retriever = vector_db.as_retriever(search_kwargs={"k": 3})

system_prompt = "現在開始使用我提供的情境來回答,只能使用繁體中文,不要有簡體中文字。如果你不確定答案,就說不知道。情境如下:\n\n{context}"
prompt_template = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("user", "問題: {input}"),
]
)

我們前面還在自己手刻 calc_similar_vectors() 函數來計算向量的相似度,這裡直接幫我們都處理好了,而且我相信一定比我們自己刻的還厲害。同時這裡使用 ChatPromptTemplate 來建立樣板,如果仔細看內容應該會覺得不算太陌生。

最後把所有的東西組裝起來,就可以來聊天了:

檔案:rag_v2.py
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

# ... 略 ...

document_chain = create_stuff_documents_chain(llm, prompt_template)
retrieval_chain = create_retrieval_chain(retriever, document_chain)

context = []
input_text = input("您想問什麼問題?\n>>> ")

while input_text.lower() != "bye":
response = retrieval_chain.invoke({"input": input_text, "context": context})
context = response["context"]

print(response["answer"])

input_text = input(">>> ")

執行看看:

python rag_v2.py
您想問什麼問題?
>>> 高見龍是誰?
回答:高見龍是一位愛寫程式的電腦阿宅,並且希望可以寫一輩子程式。他是父母給他的本名,這個名字看起來有點像武俠小說的名字,但其實不是筆名。他很喜歡自己的名字。
>>> 他有出版過什麼書?
回答: 不知道。

雖然不知道出過什麼書,不知道是還不夠聰明還是提供的資料還不夠所以無法判斷出過什麼書,但還是比較亂講一通好多了。

完整程式碼如下:

檔案:rag_v2.py
from langchain_community.llms import Ollama
from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

llm = Ollama(model="mistral")

loader = TextLoader("about_me.txt")

text_splitter = RecursiveCharacterTextSplitter(
chunk_size=20,
chunk_overlap=5,
separators=[" ", ",", "\n"],
)

splited_docs = text_splitter.split_documents(loader.load())

embeddings = OllamaEmbeddings(model="nomic-embed-text")

vector_db = Chroma.from_documents(
documents=splited_docs,
embedding=embeddings,
persist_directory="db",
collection_name="about",
)

retriever = vector_db.as_retriever(search_kwargs={"k": 3})

system_prompt = "現在開始使用我提供的情境來回答,只能使用繁體中文,不要有簡體中文字。如果你不確定答案,就說不知道。情境如下:\n\n{context}"
prompt_template = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("user", "問題: {input}"),
]
)

document_chain = create_stuff_documents_chain(llm, prompt_template)
retrieval_chain = create_retrieval_chain(retriever, document_chain)

context = []
input_text = input("您想問什麼問題?\n>>> ")

while input_text.lower() != "bye":
response = retrieval_chain.invoke({"input": input_text, "context": context})
context = response["context"]

print(response["answer"])

input_text = input(">>> ")

讀取 PDF 檔案

剛才都是處理文字檔,但其實要處理 PDF 檔案也不難,只要換個 Loader 就可以了:

檔案:rag_v3.py
from langchain_community.document_loaders import PyPDFLoader

# ... 略 ...

loader = PyPDFLoader("interview.pdf")
splited_docs = loader.load_and_split()

這裡我準備了一個 interview.pdf 的文件檔,並改用 PyPDFLoader 來讀取 PDF 檔案,並且用 load_and_split() 方法來把 PDF 檔案的內容分割成小段落。另外因為不想讓向資都存在同一個 collection 裡,所以我把 collection_name 換成 interview

檔案:rag_v3.py
# ... 略 ...
vector_db = Chroma.from_documents(
documents=splited_docs,
embedding=embeddings,
persist_directory="db",
collection_name="interview",
)

執行看看:

$ python rag_v3.py
您想問什麼問題?
>>> 什麼是快速面試?
快速面試是一種幫我們的培訓班學員找新工程師工作的媒合活動,旨在幫助企業找到合適的工程師,同時又能夠讓學員們找到適合的新工作機會。這個活動可線上或是實體進行,每家企業與每位求職者進行8分鐘的面試,活動結束後我們會提供通過快速面試的求職者資訊給企業。這個活動是免費進行的,需要準備相關公司和職缺資料以及對面試者的能力進行有效的評估。
>>> 下次什麼時候舉辦?
答案:下一場快速⾯試預計在2024年10月2日(星期三)舉辦。

看起來不錯!

完整程式碼:

檔案:rag_v3.py
from langchain_community.llms import Ollama
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.embeddings import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains import create_retrieval_chain

llm = Ollama(model="mistral")

loader = PyPDFLoader("interview.pdf")
splited_docs = loader.load_and_split()

embeddings = OllamaEmbeddings(model="nomic-embed-text")

vector_db = Chroma.from_documents(
documents=splited_docs,
embedding=embeddings,
persist_directory="db",
collection_name="interview",
)

retriever = vector_db.as_retriever(search_kwargs={"k": 3})

system_prompt = "現在開始使用我提供的情境來回答,只能使用繁體中文,不要有簡體中文字。如果你不確定答案,就說不知道。情境如下:\n\n{context}"
prompt_template = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("user", "問題: {input}"),
]
)

document_chain = create_stuff_documents_chain(llm, prompt_template)
retrieval_chain = create_retrieval_chain(retriever, document_chain)

context = []
input_text = input("您想問什麼問題?\n>>> ")

while input_text.lower() != "bye":
response = retrieval_chain.invoke({"input": input_text, "context": context})
context = response["context"]

print(response["answer"])

input_text = input(">>> ")