Ollama using Chinese embedding model and Qwen to answer question

sample.py

import ollama
import chromadb
import os

# Read files in `data` dir as string
# and create a list of documents
data_dir = 'data'
files = os.listdir(data_dir)
documents = []
for file in files:
    with open(os.path.join(data_dir, file), 'r') as f:
        documents.append(f.read())

# Initialize ChromaDB
chroma_client = chromadb.PersistentClient(path='./chroma.db')

# Recreate collection if exists to keep it clean
try:
    chroma_client.get_collection("roasted_pig")
    chroma_client.delete_collection("roasted_pig")
    collection = chroma_client.create_collection(name="roasted_pig")
except ValueError as e:
    print("Collection does not exist, creating collection...")
    collection = chroma_client.create_collection(name="roasted_pig")

# Embed document and add it ChromaDB
for i, doc in enumerate(documents):
    # Chinese embedding model: https://ollama.com/herald/dmeta-embedding-zh
    # ollama pull herald/dmeta-embedding-zh
    response = ollama.embeddings(model='herald/dmeta-embedding-zh', prompt=doc)
    embedding = response['embedding']
    collection.add(
        ids=[str(i)],
        documents=[doc],
        embeddings=[embedding]
    )

# Retrieve data - Embed the question and search in collection
prompt = '我想買回門乳豬'
response = ollama.embeddings(model='herald/dmeta-embedding-zh', prompt=prompt)
result = collection.query(
    query_embeddings=[response['embedding']], n_results=1
)

# print(result['documents'][0][0])
# Put the result to LLM (Qwen - 通義千問1.5:4B (default)) to generate response
system_msg = f"你是一位對親切有禮的客戶服務人員。你僅能用這些資訊來回答你的問題。若你不知道，請回答：我不知道。你僅能使用這些資訊： {result['documents'][0][0]} 來回答你的問題。不可以產生資訊之外的文字。"
prompt = f"請回答客人這個問題： {prompt} "

# If use ollama.generate(), use output['response'].  output is `dict`
# output = ollama.generate(model='qwen', prompt = prompt, system=system_msg)
# print(output['response'])

# If use ollama.chat(), use output['message']['content'].  output is `dict`
output = ollama.chat(model='qwen', 
                     messages=[
                       {'role': 'system', 'content': system_msg}, 
                       {'role': 'user', 'content': prompt}
                     ])
print(output['message']['content'])