Merge pull request #42 from JustInternetAI/Chloe

Chloe

Author: chllein

.docker/Dockerfile

@@ -13,7 +13,7 @@ RUN apt-get update && \
 # Install requirements
 COPY requirements.txt /workspace/requirements.txt
 RUN pip install --no-cache-dir -r /workspace/requirements.txt
-RUN pip install torch==2.1.2 torchvision --index-url https://download.pytorch.org/whl/cpu
+RUN pip install torch==2.4.0 torchvision --index-url https://download.pytorch.org/whl/cpu
 
 
 RUN playwright install

src/justinsight/celery.py

@@ -23,11 +23,11 @@
     #     "args": (),
     # },
 
-    "check-BBCfeed-every-5-minutes": {
-        "task": "justinsight.tasks.bbcLogger_task",
-        "schedule": 5.0,
-        "args": (),
-    },
+    # "check-BBCfeed-every-5-minutes": {
+    #     "task": "justinsight.tasks.bbcLogger_task",
+    #     "schedule": 5.0,
+    #     "args": (),
+    # },
 
     # "check-CBSfeed-every-5-minutes": {
     #     "task": "justinsight.tasks.cbsLogger_task",

src/justinsight/tasks.py

@@ -10,7 +10,7 @@
 from ingest.usnews_ingestor import USNEWSIngestor
 from .nlpthings import dummy_addToEntryInDB
 from ingest.save_to_database import collection
-from nlp.core import process_article
+from nlp.ner_core import NERCore
 from bson import ObjectId
 
 @shared_task
@@ -90,8 +90,10 @@ def runNER_task(entry_id):
 @shared_task
 def ner_task(article_id):
     # Process article with NER results
+    print("In the NER task")
+    core = NERCore()
     print("Actually trying to do NER!!!")
-    process_article(article_id)
+    core.process_article(article_id)
 
 
 

src/nlp/base_core.py

@@ -0,0 +1,53 @@
+from pymongo import MongoClient
+from bson import ObjectId
+from transformers import pipeline
+
+class BaseCore:
+    # Include username, password, and authentication database
+    client = MongoClient("mongodb://myuser:mypassword@mongo:27017/justinsightdb?authSource=admin")
+    db = client["justinsightdb"]
+    collection = db["articles"]
+
+    model = None #Set in subclass
+
+    def __init__(self, task: str, model_name: str, aggregation_strategy: str = None):
+        self.task = task
+        self.model_name = model_name
+        self.aggregation_strategy = aggregation_strategy
+        self.pipeline = self.load_pipeline()
+
+    def load_pipeline(self):
+        if self.task == "ner" and self.aggregation_strategy:
+            return pipeline(self.task, model=self.model_name, aggregation_strategy=self.aggregation_strategy)
+
+        print(f"pipeline loading...")
+        return pipeline(self.task, model=self.model_name)
+
+    def process_article(self, article_id: str):
+        #Retrieve article by ID
+        article = self.collection.find_one({"_id": ObjectId(article_id)})
+
+        if not article:
+            print(f"No article found with ID: {article_id}")
+            return []
+
+        if article.get("processed") is True:
+            print(f"Article {article_id} already processed.")
+            return
+
+        # We have a check running so only articles with full text are saved
+        full_text = article.get("full_text", "")
+        # if not full_text:
+        #     print(f"Article {article_id} has no full text.")
+        #     return
+
+        entities = self.pipeline(full_text)
+
+        self.collection.update_one(
+            {"_id": ObjectId(article_id)},
+            {"$set": {
+                "processed": True,
+                "entities": entities
+            }}
+        )
+        return entities

src/nlp/core.py

@@ -0,0 +1,64 @@
+from nlp.base_core import BaseCore
+from transformers import pipeline
+from bson import ObjectId
+
+class NERCore(BaseCore):
+    def __init__(self):
+        print("constructing NER Core instance")
+        super().__init__(
+            task="ner",
+            model_name="dslim/bert-base-NER",
+            aggregation_strategy="simple"
+        )
+    
+    def process_article(self, article_id: str):
+        #Retrieve article by ID
+        article = self.collection.find_one({"_id": ObjectId(article_id)})
+
+        if not article:
+            print(f"No article found with ID: {article_id}")
+            return []
+
+        if article.get("processed") is True:
+            print(f"Article {article_id} already processed.")
+            return []
+
+        # We have a check running so only articles with full text are saved
+        full_text = article.get("full_text", "")
+        # if not full_text:
+        #     print(f"Article {article_id} has no full text.")
+        #     return
+
+        # Run NER
+        entities = self.pipeline(full_text) # run_ner_hf(full_text)
+ 
+        # Update article in DB
+        self.addToEntryInDB(article_id, {
+            "ner": entities,
+            "processed": True
+        })
+
+        return entities
+
+    def format_ner_tags(self, ner_list):
+        formatted = []
+        for ent in ner_list:
+            label = ent.get("label") or ent.get("entity") or ent.get("entity_group", "UNKNOWN")
+            text = ent.get("text") or ent.get("word") or ""
+            formatted.append(f"{label}: {text}")
+        return ", ".join(formatted)
+
+    def addToEntryInDB(self, entry_id, updates):
+        print("Adding NER results to database\r\r\r")
+
+        if "ner" in updates:
+            for ent in updates["ner"]:
+                ent["score"] = float(ent["score"])  # convert np.float32 to Python float
+                
+            updates["ner_pretty"] = self.format_ner_tags(updates["ner"]) # so we can actually read the NER
+
+        id = ObjectId(entry_id)
+        self.collection.update_one(
+            {"_id": id},
+            {"$set": updates}
+        )