StockDocs/embedding/embedder.py

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)