Spaces:
Sleeping
Sleeping
File size: 1,920 Bytes
31f0f40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | import re
import os
import duckdb
import pandas as pd
from .embeddings import model
def split_into_sentences(text):
text = re.sub(r'\s+', ' ', text).strip()
fragments = text.split('. ')
return [f.strip() for f in fragments if f.strip()]
def init_document_db(csv_path="data/case-teaching-cabot.csv"):
db_path = "document.db"
if os.path.exists(db_path):
os.remove(db_path)
con = duckdb.connect(db_path)
con.execute("CREATE SEQUENCE document_seq START 1;")
con.execute("CREATE SEQUENCE sentence_seq START 1;")
con.execute("""
CREATE TABLE document (
document_id INTEGER PRIMARY KEY DEFAULT nextval('document_seq'),
content TEXT
);
""")
con.execute("""
CREATE TABLE sentence (
sentence_id INTEGER PRIMARY KEY DEFAULT nextval('sentence_seq'),
document_id INTEGER REFERENCES document(document_id),
content TEXT
);
""")
con.execute("""
CREATE TABLE sentence_embedding (
sentence_id INTEGER REFERENCES sentence(sentence_id),
embedding DOUBLE[]
);
""")
df = pd.read_csv(csv_path)
for _, row in df.iterrows():
doc_text = str(row['document'])
document_id = con.execute(
"INSERT INTO document (content) VALUES (?) RETURNING document_id",
[doc_text]
).fetchone()[0]
for sentence_text in split_into_sentences(doc_text):
con.execute(
"INSERT INTO sentence (document_id, content) VALUES (?, ?)",
[document_id, sentence_text]
)
sentences = con.execute("SELECT sentence_id, content FROM sentence").fetchall()
for sentence_id, content in sentences:
embedding = model.encode(content)
con.execute(
"INSERT INTO sentence_embedding VALUES (?, ?)",
[sentence_id, embedding.tolist()]
)
return con
|