Add keyword extraction logic and rename source text collector directory

Tests/test_html_tag_collector_integration.py → Tests/test_source_text_collector_integration.py

@@ -1,7 +1,7 @@
 import polars as pl
 import pytest
 
-from html_tag_collector.collector import process_in_batches
+from source_text_collector.collector import process_in_batches
 
 sample_json_data = [{
     "id": 1,

Tests/test_source_text_collector_unit.py

@@ -0,0 +1,16 @@
+import pytest
+from unittest.mock import MagicMock, patch
+from source_text_collector.keyword_extractor import KeywordExtractor
+
+@patch('keybert.KeyBERT.extract_keywords')
+def test_extract_keywords(mock_extract_keywords):
+    # Arrange
+    mock_extract_keywords.return_value = [('keyword1', 0.8), ('keyword2', 0.7), ('keyword3', 0.6)]  # return some mock keywords
+    keyword_extractor = KeywordExtractor()
+
+    # Act
+    keywords = keyword_extractor.extract_keywords("some text", n=3)
+
+    # Assert
+    mock_extract_keywords.assert_called_once_with("some text", keyphrase_ngram_range=(1, 3), stop_words='english', top_n=3)
+    assert keywords == ['keyword1', 'keyword2', 'keyword3']

identification_pipeline.py

@@ -2,7 +2,7 @@
 import os
 import polars as pl
 import sys
-from html_tag_collector.collector import collector_main
+from source_text_collector.collector import collector_main
 from agency_identifier.identifier import match_urls_to_agencies_and_clean_data
 from datetime import datetime as dt
 from dotenv import load_dotenv

html_tag_collector/collector.py → source_text_collector/collector.py

@@ -18,12 +18,14 @@
 import polars as pl
 
 from RootURLCache import RootURLCache
+from keyword_extractor import KeywordExtractor
 
 # Define the list of header tags we want to extract
 header_tags = ["h1", "h2", "h3", "h4", "h5", "h6"]
 DEBUG = False  # Set to True to enable debug output
 VERBOSE = False  # Set to True to print dataframe each batch
 root_url_cache = RootURLCache()
+keyword_extractor = KeywordExtractor()
 
 
 def process_urls(manager_list, render_javascript):
@@ -278,6 +280,11 @@ def parse_response(url_response):
         tags["html_title"] = ""
     tags["root_page_title"] = root_url_cache.get_title(tags["url"])
 
+    tags["keywords"] = str(keyword_extractor.extract_keywords(
+        text=soup.get_text(separator=' ', strip=True),
+        n=10
+    ))
+
     meta_tag = soup.find("meta", attrs={"name": "description"})
     try:
         tags["meta_description"] = meta_tag["content"] if meta_tag is not None else ""

source_text_collector/keyword_extractor.py

@@ -0,0 +1,29 @@
+import asyncio
+
+import requests
+import yake
+from bs4 import BeautifulSoup
+from keybert import KeyBERT
+
+
+# from multi_rake import Rake
+
+class KeywordExtractor:
+    """
+    This class is used to extract keywords from text using the KeyBERT algorithm.
+    """
+
+    def __init__(self):
+        self.model = KeyBERT(model='all-mpnet-base-v2')
+
+    def extract_keywords(self, text: str, n=10) -> list[str]:
+        """
+        Extract keywords from text using KeyBERT algorithm.
+        Args:
+            text: The text to extract keywords from.
+            n: The number of keywords to extract.
+
+        Returns: The top n keywords extracted from the text.
+        """
+        keywords = self.model.extract_keywords(text, keyphrase_ngram_range=(1, 3), stop_words='english', top_n=n)
+        return [keyword for keyword, _ in keywords]

tests/test_root_url_cache_unit.py

@@ -3,7 +3,7 @@
 import os
 import pytest
 from unittest.mock import mock_open, patch
-from html_tag_collector.RootURLCache import RootURLCache  # Adjust import according to your package structure
+from source_text_collector.RootURLCache import RootURLCache  # Adjust import according to your package structure
 
 
 @pytest.fixture