Spaces:
Running
Running
File size: 1,792 Bytes
4116826 |
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 |
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
import os
import pandas as pd
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Load CVE dataset
df = pd.read_csv("cve.csv")
# Set up the embedding model using HuggingFace with a fully qualified model name
# Using a simpler model that's more compatible with Hugging Face Spaces
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/paraphrase-MiniLM-L3-v2", # Smaller, more compatible model
model_kwargs={'device': 'cpu'} # Ensure it runs on CPU for compatibility
)
# Directory for the vector store - use /tmp for proper permissions in containerized environments
db_location = "/tmp/chrome_langchain_db"
add_documents = not os.path.exists(db_location)
# Initialize Chroma DB
vector_store = Chroma(
collection_name="cve_data",
persist_directory=db_location,
embedding_function=embeddings
)
# Add documents only if DB doesn't exist yet
if add_documents:
documents = []
ids = []
for i, row in df.iterrows():
# Replace with actual column names in your CSV
cve_id = row.get("CVE_ID", f"CVE-{i}")
description = row.get("Description", "")
date = row.get("PublishedDate", "")
content = f"CVE ID: {cve_id}\nDescription: {description}\nPublished Date: {date}"
document = Document(
page_content=content,
metadata={"published_date": date},
id=str(i)
)
documents.append(document)
ids.append(str(i))
vector_store.add_documents(documents=documents, ids=ids)
# Create retriever from the vector store
retriever = vector_store.as_retriever(search_kwargs={"k": 5})
|