91 lines
3.7 KiB
Python
91 lines
3.7 KiB
Python
import os
|
|
import json
|
|
import chromadb
|
|
import uuid
|
|
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
model = SentenceTransformer(
|
|
"Snowflake/snowflake-arctic-embed-m-long",
|
|
device="cpu", # <-- the only line that changes
|
|
trust_remote_code=True
|
|
)
|
|
|
|
|
|
import math
|
|
def embed_text(text, specific_context, max_tokens=2048, overlap=256):
|
|
"""
|
|
Embeds the given text using the SentenceTransformer model.
|
|
Splits the text into chunks if it exceeds max_tokens.
|
|
"""
|
|
if len(text) + len(specific_context) <= max_tokens:
|
|
return (text + specific_context, model.encode([text, specific_context], normalize_embeddings=True))
|
|
|
|
# Split text into chunks
|
|
chunks = []
|
|
start = 0
|
|
|
|
chunkNum = math.ceil(len(text) / (max_tokens - len(specific_context)))
|
|
# 1000 + 20 = 1020
|
|
# 3000 + 20 = 3020 -> 3000 / (2048 - 20) = 1.47 ~ 2
|
|
|
|
for i in range(chunkNum):
|
|
if start >= len(text):
|
|
break
|
|
# Calculate end index for the chunk
|
|
end = min(start + max_tokens - len(specific_context), len(text))
|
|
chunk = text[start:end] + specific_context
|
|
chunks.append(chunk)
|
|
start += (max_tokens - overlap) # Overlap for next chunk
|
|
|
|
return (chunks, model.encode(chunks, normalize_embeddings=True))
|
|
|
|
|
|
client = chromadb.HttpClient(host="localhost", port=8000)
|
|
|
|
# client.delete_collection("news") # Replace "news" with your collection name
|
|
|
|
collection = client.get_or_create_collection("news")
|
|
|
|
# Load all the processed articles from the output folder
|
|
output_folder = os.path.join(os.path.dirname(__file__), "../ai_processor/output")
|
|
output_articles = []
|
|
# When embedding and upserting, use the source from the processed result
|
|
for filename in os.listdir(output_folder):
|
|
if filename.endswith(".json"):
|
|
with open(os.path.join(output_folder, filename), 'r', encoding='utf-8') as f:
|
|
prev_proc = json.load(f)
|
|
output_articles.append(prev_proc)
|
|
|
|
def makeTickerMovement(ticker, movement, percent_change):
|
|
"""
|
|
Helper function to create a ticker movement entry.
|
|
"""
|
|
ticker = "" if ticker is None else ticker
|
|
movement = "" if movement is None else movement
|
|
percent_change = "" if percent_change is None else percent_change
|
|
return [ticker, movement, percent_change]
|
|
|
|
# Embed the article content
|
|
try:
|
|
specific_context = prev_proc['summary'] + ",".join([' '.join(makeTickerMovement(*x)) for x in prev_proc['tickersAndMovements']])
|
|
chunks, embedded_content = embed_text(prev_proc['original_content'], prev_proc['summary'])
|
|
|
|
print(f"Embedded content for {filename}: \n {embedded_content[:10]}...") # Print first 10 values for preview
|
|
|
|
collection.upsert(
|
|
ids = [ str(uuid.uuid4()) for d in chunks ], # Generate unique IDs for each embedded content
|
|
documents = chunks, # optional but nice for debugging
|
|
embeddings = embedded_content.tolist(), # Chroma expects List[List[float]]
|
|
metadatas =[{
|
|
"source": prev_proc.get('source', 'Unknown'), # Use the source from the processed result
|
|
"published": prev_proc.get('published', 'Unknown'),
|
|
"filename": prev_proc.get('filename', 'Unknown') # Add filename for traceability
|
|
} for d in chunks], # Metadata for each embedded content
|
|
)
|
|
|
|
except Exception as e:
|
|
print(f"Error embedding content for {filename}: {e}")
|
|
continue
|
|
|