ingest.py

from docling.document_converter import DocumentConverter
import re
import os
import unicodedata
import ollama
from supabase import create_client
from dotenv import load_dotenv

def clean_text(text:str) -> str:
    text = text.replace('-\n', '') # 删除连字符空行
    text = text.replace('\r\n', '\n') # 统一换行符号
    text = re.sub(' +',' ', text) # re.sub('pattern', 'replacement', text) '+‘ 代表一个或多个
    text = re.sub('\n{3,}', '\n\n', text) # 多个空行变成一个空行

    cleaned = []
    for c in text:
        if unicodedata.category(c).startswith('C'): # Unicodedata 所有垃圾字符都以C开头
            continue
        cleaned.append(c)
    text = ''.join(cleaned) # text 本身包含' '

    return text


# Chuck,  把书拆成多个向量，embedding 模型能处理而不截断，信息更专注 （如果embedding 整本书，向量和qury 的距离很难近）
def chuck_text(text, size, overlap):
    chucked = []
    step = size - overlap
    for i in range(0, len(text), step):  # 0开始，每次跳step 个字符
        chucked.append(text[i:i+size])  # 不是切step 个，是size 个（包含overlap)
    return chucked

# 逐个处理chucks
# def embed_chucks(chucks):
#     embeddings = []
#     for i in chucks: 
#         response = ollama.embeddings(model = 'qwen3-embedding:4b', prompt=i) # ollama.embedding(*, model, prompt)
#         embedding = response['embedding'] # response 返回字典 {'embedding': []}
#         embeddings.append(embedding)
#     return embeddings

# 一次请求处理所有chuck
def embed_chucks(chunks):
    response = ollama.embed(model='qwen3-embedding:4b', input=chunks)
    return response['embeddings']


def upload_to_supabase(chunks, embeddings, book_title, supabase):
    rows = []
    for chunk, embedding in zip(chunks, embeddings): 
        rows.append({
        'book_title':book_title,
        'content': chunk,
        'embedding': embedding,
                                })
        
    supabase.table('book').insert(rows).execute()

def main():
    load_dotenv() # 把.env 文件里的内容加载到环境变量里

    SUPABASE_URL = os.getenv('SUPABASE_URL') # 从环境变量里读取
    SUPABASE_KEY = os.getenv('SUPABASE_KEY')

    book_title = 'Cognitive_Behavior_Therapy_English'
    pdf_path = "Cognitive_Behavior_Therapy_English.pdf"

    supabase = create_client(SUPABASE_URL, SUPABASE_KEY)

    converter = DocumentConverter()
    result = converter.convert(pdf_path)
    raw_text = result.document.export_to_markdown()


    full_text = clean_text(raw_text)
    chunks = chuck_text(full_text, 3000, overlap=400)
    embeddings = embed_chucks(chunks)

    upload_to_supabase(chunks, embeddings, book_title, supabase)

if __name__ == '__main__':
    main()