136 lines
5.4 KiB
Python
136 lines
5.4 KiB
Python
import os
|
|
import json
|
|
import chromadb
|
|
import uuid
|
|
import time
|
|
import datetime
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
CHROMADB_HOST = os.getenv("CHROMADB_HOST")
|
|
CHROMADB_PORT = os.getenv("CHROMADB_PORT")
|
|
|
|
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=CHROMADB_HOST, port=CHROMADB_PORT)
|
|
|
|
# client.delete_collection("news") # Replace "news" with your collection name
|
|
|
|
collection = client.get_or_create_collection("news")
|
|
|
|
# Load cache of previously embedded articles
|
|
cache_file = "embedded_articles_cache.json"
|
|
|
|
# Load all the processed articles from the output folder
|
|
output_folder = "/app/output"
|
|
|
|
while True:
|
|
print(f"Starting embedding run at {datetime.datetime.now()}")
|
|
|
|
embedded_cache = {}
|
|
if os.path.exists(cache_file):
|
|
with open(cache_file, 'r', encoding='utf-8') as f:
|
|
embedded_cache = json.load(f)
|
|
|
|
output_articles = []
|
|
processed_articles = []
|
|
# When embedding and upserting, use the source from the processed result
|
|
for filename in os.listdir(output_folder):
|
|
if filename.endswith(".json"):
|
|
# Check if already embedded
|
|
if filename in embedded_cache:
|
|
print(f"Article {filename} already embedded, skipping.")
|
|
continue
|
|
|
|
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
|
|
)
|
|
|
|
# Add to processed list for cache update
|
|
processed_articles.append(filename)
|
|
print(f"Successfully embedded {filename}")
|
|
|
|
except Exception as e:
|
|
print(f"Error embedding content for {filename}: {e}")
|
|
continue
|
|
|
|
# Update cache with newly processed articles
|
|
for filename in processed_articles:
|
|
embedded_cache[filename] = {
|
|
"embedded_date": str(datetime.datetime.now()),
|
|
"status": "completed"
|
|
}
|
|
|
|
# Save updated cache
|
|
try:
|
|
with open(cache_file, 'w', encoding='utf-8') as f:
|
|
json.dump(embedded_cache, f, indent=2)
|
|
print(f"Updated cache with {len(processed_articles)} newly embedded articles.")
|
|
except Exception as e:
|
|
print(f"Error saving cache file: {e}")
|
|
|
|
print(f"Embedding process completed. Total articles in cache: {len(embedded_cache)}")
|
|
print("Waiting 5 minutes before next run...")
|
|
time.sleep(300) # Wait 5 minutes (300 seconds)
|
|
|