DIY Chat with Document with Ollama

Trong bài này, chúng ta sẽ build một chương trình Chat with Document đơn giản cho phép người dùng sử dụng Ollama để hỏi đáp với AI dựa trên context là một tài liệu bất kì, và mọi thứ sẽ diễn ra offline. Trong phạm vi bài viết, chúng ta sẽ giới hạn với việc chỉ load file text (ví dụ markdown), tuy nhiên, chúng ta có thể áp dụng cách làm tương tự đối với các nội dung phức tạp hơn như file PDF hoặc toàn bộ codebase của một dự án code bất kì.

Toàn bộ code trong bài viết này chỉ là pseudo code, minh hoạ ý tưởng là chính, các bạn có thể vào link https://github.com/huytd/txtask để xem implement thực tế.

Mô tả cách hoạt động

Cách hoạt động của chương trình này được chia làm 2 giai đoạn:

• Giai đoạn indexing
• Giai đoạn semantic search và hỏi đáp

Trước khi bắt đầu, thì cần nói sơ về khái niệm Embedding: đây có thể hiểu đơn giản là cách mà các language models mã hoá các thông tin đầu vào (hình ảnh, âm thanh, văn bản,...) dưới dạng một mảng (vector) gồm hàng nghìn con số khác nhau. Dựa trên các vector này, các language models có thể "hiểu" được nội dung, ngữ cảnh, ý nghĩa của từng input.

Ở giai đoạn indexing, chúng ta sẽ chia các file markdown thành từng đoạn văn bản nhỏ hơn và dùng Ollama để tính ra embedding của từng đoạn đó. Sau đó lưu tất cả các cặp đoạn văn và embedding này lại trong một vector store nào đó (trong bài này chúng ta sẽ tự implement một in-memory vector store đơn giản).

Tiếp theo, ở giai đoạn hỏi đáp, chương trình sẽ nhận vào câu hỏi, nhập bởi người dùng. Sau đó sử dụng Ollama để tính ra embedding tương ứng. Chúng ta sẽ so sánh vector embedding của câu hỏi này với từng bản ghi trong vector store đã tạo ở bước indexing, tìm ra các vector (và các đoạn nội dung tương ứng) gần với vector của câu hỏi nhất. Sử dụng các nội dung này để cung cấp thông tin, và gửi kèm câu hỏi đến cho Ollama để trả lời.

Có rất nhiều cách để tính toán độ tương đồng giữa các vector như dùng kd-tree, hoặc tính toán bằng công thức khoảng cách Euclid, cosine similarity hay tích vô hướng (dot product). Có thể đọc thêm về các cách tính này tại link: https://www.pinecone.io/learn/vector-similarity/

Trong phạm vi bài này, chúng ta sẽ tính độ tương đồng bằng cosine similarity.

Bước 1: Xây dựng vector store

Đầu tiên thì chúng ta sẽ xây dựng một chiếc vector store, chứa dữ liệu trực tiếp trên RAM. Vì sao? Vì để cho đơn giản, và điều này cũng có nghĩa là, toàn bộ dữ liệu chứa trong đó sẽ được siêu thoát khi nhỡ may chương trình crash. Nhưng được cái là nó đơn giản, dễ làm.

Vector Store của chúng ta sẽ có 2 chức năng chính:

• Tiếp nhận một đoạn văn bản (tạm gọi là S), gọi Ollama API để lấy về embedding (tạm gọi là E) cho đoạn nội dung đó, và lưu cặp (S, E) vào bộ nhớ.
• Nhận vào một đoạn văn bản (Q), gọi Ollama API lấy ra embedding của nó (Q_e), rồi dùng cosine similarity để search các cặp (S, E) có trong bộ nhớ, trả về danh sách các cặp có E tương đồng với Q_e.

Kiểu dữ liệu cho embedding thực chất là một vector của các con số, nên chúng ta có thể dùng kiểu Vec<f64>, mã giả cho vector store đại khái như này:

type Embedding = Vec<f64>;

struct VectorStore {
	data: Vec<(String, Embedding)>
}

impl VectorStore {
	fn insert_document(&mut self, content: &str) {
		embedding = [get the embedding for content];
		push the pair (content, embedding) to data;
	}

	fn find_similar(&self, input: &str) -> Vec<String> {
		Qe = [get the embedding for input];
		init result as an empty string vector;
		for each (content, embedding) in self.data {
			distance = cosine similarity between (embedding and Qe);
			if [distance is in acceptable range] {
				push content into result;
			}
		}
		return result;
	}
}

Để generate embedding cho một đoạn text, chúng ta gọi API endpoint /api/embeddings của Ollama, ví dụ:

# Request:

curl http://localhost:11434/api/embeddings -d '{
  "model": "llama2",
  "prompt": "Here is an article about llamas..."
}'

# Response:

{
  "embedding": [
    0.5670403838157654, 0.009260174818336964,...
  ]
}

Xong rồi, giờ chúng ta có thể bắt tay vào thực hiện giai đoạn 1: indexing.

Bước 2: Indexing

Ở bước này, chúng ta load toàn bộ nội dung các file input markdown, tiến hành chia nhỏ thành từng đoạn và dùng hàm insert_document của VectorStore để ghi vào database

Lưu ý là, nên chia văn bản thành nhiều đoạn văn nhỏ nhất có thể, mỗi đoạn nên chứa đầy đủ nội dung, và có nghĩa. Việc này giúp bước semantic search diễn ra chính xác hơn.

Chúng ta có thể split một file thành từng line riêng biệt để index, hoặc có thể dùng một language model nào đó cho việc này.

Nếu dùng Rust, bạn có thể dùng crate text-splitter để thực hiện việc split này.

database = new VectorStore();

for each file in data folder {
	initialize text splitter;
	chunks = splitter.split(file);
	for each chunk in chunks {
		database.insert_document(chunk);
	}
}

Bước 3: Semantic search và hỏi đáp

Ở bước này, chúng ta tiến hành nhận vào câu hỏi từ user, sử dụng hàm find_similar của VectorStore, lấy ra danh sách các đoạn nội dung có liên quan đến câu hỏi.

Ví dụ, nếu trong các file markdown của mình có một số đoạn nội dung như:

Vinasat-1 was launched on April 18, 2008, and was built by Lockheed Martin for Vietnam Posts and Telecommunications Group (VNPT).

The Moon Landing Program refers to the historical efforts of various countries to land humans on the Moon.

The most well-known of these programs is the Apollo program by NASA, which successfully landed humans on the Moon in 1969 4.

Và câu hỏi được nhập vào là: "When did we first launch Vinasat", thì đoạn văn nói về Vinasat sẽ có điểm tương đồng cao nhất, và sẽ nằm trong danh sách các string được trả về bởi hàm find_similar.

Sau khi có danh sách các string tương ứng, chúng ta tiến hành viết một prompt kiểu như này:

Answer my question based on the provided context.

Context:

{list of all matching string}

Question:

{the input question}

Và dùng prompt này để gọi API endpoint /api/chat của Ollama và lấy về câu trả lời:

question = read string from stdin;
matching_strings = database.find_similar(question);

prompt = "Answer my question based on the provided context.\nContext:"

for each content in matching_strings {
	prompt += content;
}

ollama_payload = {
  "messages": [
    {
      "role": "user",
      "content": prompt
    }
  ]
}

POST ollama_payload to http://localhost:11434/api/chat

Đến đây thì chúng ta đã hoàn thành chương trình hỏi đáp với văn bản. Việc tiếp theo là gắn thêm Web UI vào, add thêm chức năng login, thanh toán hoặc subscription, làm landing page, viết document hướng dẫn sử dụng, kiếm server để deploy Ollama lên, hoặc thay Ollama thành OpenAI API, quay video demo sản phẩm để post lên Twitter, launch lên ProductHunt và đem đi bán để cạnh tranh với các sản phẩm Chat to PDF/Document nhan nhản khắp nơi thôi.