Spaces:
Sleeping
Sleeping
File size: 6,584 Bytes
cf10c85 |
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 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
import os
import feedparser
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document
import logging
from huggingface_hub import HfApi, login, snapshot_download
import shutil
import json
from datetime import datetime
import dateutil.parser
import hashlib
import re
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
LOCAL_DB_DIR = "chroma_db"
COLLECTION_NAME = "news_articles"
HF_API_TOKEN = os.getenv("DEMO_HF_API_TOKEN", "YOUR_HF_API_TOKEN")
REPO_ID = "broadfield-dev/news-rag-db"
FEEDS_FILE = "rss_feeds.json"
login(token=HF_API_TOKEN)
hf_api = HfApi()
def get_embedding_model():
if not hasattr(get_embedding_model, "model"):
get_embedding_model.model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
return get_embedding_model.model
def clean_text(text):
if not text or not isinstance(text, str):
return ""
text = re.sub(r'<.*?>', '', text)
text = ' '.join(text.split())
return text.strip().lower()
def fetch_rss_feeds():
articles = []
seen_keys = set()
try:
with open(FEEDS_FILE, 'r') as f:
feed_categories = json.load(f)
except FileNotFoundError:
logger.error(f"{FEEDS_FILE} not found. No feeds to process.")
return []
for category, feeds in feed_categories.items():
for feed_info in feeds:
feed_url = feed_info.get("url")
if not feed_url:
continue
try:
logger.info(f"Fetching '{feed_info.get('name', feed_url)}' from category '{category}'")
# Add a User-Agent to prevent getting blocked
feed = feedparser.parse(feed_url, agent="RSSNewsBot/1.0 (+http://huggingface.co/spaces/broadfield-dev/RSS_News)")
if feed.bozo:
logger.warning(f"Parse error for {feed_url}: {feed.bozo_exception}")
continue
for entry in feed.entries[:10]: # Process max 10 entries per feed
title = entry.get("title", "No Title")
link = entry.get("link", "")
description = entry.get("summary", entry.get("description", ""))
cleaned_title = clean_text(title)
cleaned_link = clean_text(link)
published = "Unknown Date"
for date_field in ["published", "updated", "created", "pubDate"]:
if date_field in entry:
try:
parsed_date = dateutil.parser.parse(entry[date_field])
published = parsed_date.strftime("%Y-%m-%d %H:%M:%S")
break
except (ValueError, TypeError):
continue
key = f"{cleaned_title}|{cleaned_link}|{published}"
if key not in seen_keys:
seen_keys.add(key)
image = "svg"
if 'media_content' in entry and entry.media_content:
image = entry.media_content[0].get('url', 'svg')
elif 'media_thumbnail' in entry and entry.media_thumbnail:
image = entry.media_thumbnail[0].get('url', 'svg')
articles.append({
"title": title,
"link": link,
"description": description,
"published": published,
"category": category, # Directly use category from JSON
"image": image,
})
except Exception as e:
logger.error(f"Error fetching {feed_url}: {e}")
logger.info(f"Total articles fetched: {len(articles)}")
return articles
def process_and_store_articles(articles):
vector_db = Chroma(
persist_directory=LOCAL_DB_DIR,
embedding_function=get_embedding_model(),
collection_name=COLLECTION_NAME
)
try:
existing_ids = set(vector_db.get(include=[])["ids"])
except Exception:
existing_ids = set()
docs_to_add = []
ids_to_add = []
for article in articles:
cleaned_title = clean_text(article["title"])
cleaned_link = clean_text(article["link"])
doc_id = f"{cleaned_title}|{cleaned_link}|{article['published']}"
if doc_id in existing_ids:
continue
metadata = {
"title": article["title"],
"link": article["link"],
"original_description": article["description"],
"published": article["published"],
"category": article["category"],
"image": article["image"],
}
doc = Document(page_content=clean_text(article["description"]), metadata=metadata)
docs_to_add.append(doc)
ids_to_add.append(doc_id)
existing_ids.add(doc_id)
if docs_to_add:
vector_db.add_documents(documents=docs_to_add, ids=ids_to_add)
vector_db.persist()
logger.info(f"Added {len(docs_to_add)} new articles to DB. Total in DB: {vector_db._collection.count()}")
def download_from_hf_hub():
if not os.path.exists(LOCAL_DB_DIR):
try:
snapshot_download(
repo_id=REPO_ID,
repo_type="dataset",
local_dir=".",
local_dir_use_symlinks=False,
allow_patterns=f"{LOCAL_DB_DIR}/**",
token=HF_API_TOKEN
)
except Exception as e:
logger.warning(f"Could not download DB from Hub (this is normal on first run): {e}")
def upload_to_hf_hub():
if os.path.exists(LOCAL_DB_DIR):
try:
hf_api.upload_folder(
folder_path=LOCAL_DB_DIR,
path_in_repo=LOCAL_DB_DIR,
repo_id=REPO_ID,
repo_type="dataset",
token=HF_API_TOKEN,
commit_message="Update RSS news database"
)
except Exception as e:
logger.error(f"Error uploading to Hugging Face Hub: {e}")
if __name__ == "__main__":
download_from_hf_hub()
articles = fetch_rss_feeds()
if articles:
process_and_store_articles(articles)
upload_to_hf_hub() |