clean up

rag_demo/rag/source_annotator.py

@@ -11,6 +11,7 @@ from transformers import pipeline
 
 class SourceAnnotator:
     def __init__(self):
+        # Extractive question answering model
         self.source_annotator = pipeline(
             "question-answering",
             model="distilbert/distilbert-base-cased-distilled-squad",
@@ -22,7 +23,7 @@ class SourceAnnotator:
         for sentence in sentences:
             scores = []
             for chunk in reranked_chunks:
-                score = self.annotate_source(sentence.lower(), chunk.content.lower())
+                score = self.source_annotator(sentence.lower(), chunk.content.lower())
                 score["filename"] = chunk.metadata["filename"].split(".pdf")[0]
                 score["chunk_id"] = chunk.chunk_id
                 scores.append(score)
@@ -38,6 +39,3 @@ class SourceAnnotator:
         pattern = r"(?<=[.!?])\s+(?=[A-Z])"
         sentences = re.split(pattern, text)
         return [s.strip() for s in sentences if s.strip()]
-
-    def annotate_source(self, text: str, chunk: str) -> dict:
-        return self.source_annotator(text, chunk)