wehub-resource-sync
87 lines
2.2 KiB
Python
87 lines
2.2 KiB
Python
import pixeltable as pxt
|
|
from pixeltable.iterators import DocumentSplitter
|
|
from pixeltable.functions.huggingface import sentence_transformer
|
|
|
|
# Initialize app structure
|
|
pxt.drop_dir("pdf_search", force=True)
|
|
pxt.create_dir("pdf_search")
|
|
|
|
# Create documents table
|
|
documents_t = pxt.create_table(
|
|
"pdf_search.documents",
|
|
{"pdf": pxt.Document}
|
|
)
|
|
|
|
# Create chunked view for efficient processing
|
|
documents_chunks = pxt.create_view(
|
|
"pdf_search.document_chunks",
|
|
documents_t,
|
|
iterator=DocumentSplitter.create(
|
|
document=documents_t.pdf,
|
|
separators="token_limit",
|
|
limit=300 # Tokens per chunk
|
|
)
|
|
)
|
|
|
|
# Configure embedding model
|
|
embed_model = sentence_transformer.using(
|
|
model_id="intfloat/e5-large-v2"
|
|
)
|
|
|
|
# Add search capability
|
|
documents_chunks.add_embedding_index(
|
|
column="text",
|
|
string_embed=embed_model
|
|
)
|
|
|
|
# Define search query
|
|
@pxt.query
|
|
def search_documents(query_text: str, limit: int = 5):
|
|
sim = documents_chunks.text.similarity(query_text)
|
|
return (
|
|
documents_chunks.order_by(sim, asc=False)
|
|
.select(
|
|
documents_chunks.text,
|
|
similarity=sim
|
|
)
|
|
.limit(limit)
|
|
)
|
|
|
|
# Sample document URLs
|
|
DOCUMENT_URL = (
|
|
"https://github.com/pixeltable/pixeltable/raw/release/docs/resources/rag-demo/"
|
|
)
|
|
|
|
document_urls = [
|
|
DOCUMENT_URL + doc for doc in [
|
|
"Argus-Market-Digest-June-2024.pdf",
|
|
"Company-Research-Alphabet.pdf",
|
|
"Zacks-Nvidia-Report.pdf",
|
|
]
|
|
]
|
|
|
|
# Add documents to database
|
|
documents_t.insert({"pdf": url} for url in document_urls)
|
|
|
|
# Search documents
|
|
@pxt.query
|
|
def find_relevant_text(query: str, top_k: int = 5):
|
|
sim = documents_chunks.text.similarity(query)
|
|
return (
|
|
documents_chunks.order_by(sim, asc=False)
|
|
.select(
|
|
documents_chunks.text,
|
|
similarity=sim
|
|
)
|
|
.limit(top_k)
|
|
)
|
|
|
|
# Example search
|
|
results = find_relevant_text(
|
|
"What are the growth projections for tech companies?"
|
|
).collect()
|
|
|
|
# Print results
|
|
for r in results:
|
|
print(f"Similarity: {r['similarity']:.3f}")
|
|
print(f"Text: {r['text']}\n") |