rag-citation 0.0.1__tar.gz

rag_citation-0.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Rahul Anand
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

rag_citation-0.0.1/PKG-INFO

@@ -0,0 +1,205 @@
+Metadata-Version: 2.1
+Name: rag_citation
+Version: 0.0.1
+Summary: RAG Citation is an project that combines Retrieval-Augmented Generation (RAG) with automatic citation generation. This tool is designed to enhance the credibility of AI-generated content by providing relevant citations for the information used in generating responses.
+Home-page: https://github.com/rahulanand1103/rag-citation
+Author: rahul anand
+Author-email: rahulanand1103@gmail.com
+License: MIT
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: spacy==3.7.5
+Requires-Dist: sentence_transformers==3.0.1
+
+# RAG Citation: Enhancing AI-Generated Content with Automatic Citations (A Non-LLM Approach)
+
+## Project Overview
+
+RAG Citation is an project that combines Retrieval-Augmented Generation (RAG) with automatic citation generation. This tool is designed to enhance the credibility of AI-generated content by providing relevant citations for the information used in generating responses.
+
+## Key Features
+
+- **Non-LLM Approach:** Utilizes efficient algorithms and NLP techniques for citation generation, making it fast and lightweight.
+- **Semantic Search:** Identifies relevant source documents based on meaning and context rather than just keyword matching.
+- **Named Entity Recognition:**  Extracts and returns relevant named entities from LLM-generated answers, such as people, organizations, money and dates.
+- **Flexible Integration:** Can be easily integrated into rag pipeline.
+- **Hallucination (Beta)** This beta feature identifies instances where the LLM-generated answer contains entities like ["DATE", "MONEY", "CARDINAL", "ORDINAL", "QUANTITY", "TIME"], but these entities cannot be found within the context. If such a mismatch occurs, it flags the result as a potential hallucination.
+
+## Quickstart
+* <b>Langchain example</b>: [langchain.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/3.example-langchain.ipynb)
+* <b>Embeddchain example</b>: [embeddchain.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/2.example-embeddchain.ipynb)
+
+To get started with `rag-citation`, install it using pip:
+
+```bash
+pip install rag-citation
+```
+
+Here's a basic example demonstrating how to use the library:
+
+```python
+from rag_citation import CiteItem, Inference
+import uuid
+
+## Sample context from vectorDB or semantic search
+documents = [
+    "Elon MuskCEO, Tesla$221.6B$439M (0.20%)Real Time Net Worthas of 8/6/24Reflects change since 5 pm ET of prior trading day. 1 in the world todayPhoto by Martin Schoeller for ForbesAbout Elon MuskElon Musk cofounded six companies, including electric car maker Tesla, rocket producer SpaceX and tunneling startup Boring Company.He owns about 12% of Tesla excluding options, but has pledged more than half his shares as collateral for personal loans of up to $3.5 billion.In early 2024, a Delaware judge voided Musk's 2018 deal to receive options equaling an additional 9% of Tesla.",
+    "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3] Musk was born in Pretoria to model Maye and businessman and engineer Errol Musk, and briefly attended the University of Pretoria before immigrating to Canada at age 18, acquiring citizenship through his Canadian-born mother. Two years later, he matriculated at Queen's University at Kingston in Canada. Musk later transferred to the University of Pennsylvania and received bachelor's degrees in economics and physics."
+]
+
+## Example answer generated by an LLM
+answer = "Elon Musk's net worth is estimated to be US$241 billion as of August 2024."
+
+## Helper function to generate a UUID
+def generate_uuid():
+    return str(uuid.uuid4())
+
+## Helper function to create context in the correct format
+def format_document(documents):
+    context = []
+    for document in documents:
+        context.append(
+            {
+                "source_id": generate_uuid(), 
+                "document": document, 
+                "meta": [{"meta-data": "some-info"}], 
+            }
+        )
+    return context
+
+context = format_document(documents)
+cite_item = CiteItem(answer=answer, context=context)
+
+## Initialize the Inference 
+inference = Inference(spacy_model="sm", embedding_model="md")
+
+## Get citation and other information
+output = inference(cite_item)
+
+print("------ Citation ------")
+print(output.citation) 
+print("------ Hallucination ------") 
+print(output.hallucination) 
+print("------ Missing Entities ------")
+print(output.missing) 
+
+```
+
+## Output Explanation
+
+### `print(output.citation)`
+```
+[
+  {
+    "answer_sentences": "Elon Musk's net worth is estimated to be US$241 billion as of August 2024.",
+    "cite_document": [
+      {
+        "document": "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3]",
+        "source_id": "23d1f1f0-2afa-4749-8639-78ec685fd837",
+        "entity": [
+          {
+            "word": "US$241 billion",
+            "entity_name": "MONEY"
+          },
+          {
+            "word": "August 2024",
+            "entity_name": "DATE"
+          }
+        ],
+        "meta": [
+          {
+            "url": "https://www.forbes.com/profile/elon-musk/",
+            "chunk_id": "1eab8dd1ffa92906f7fc839862871ca5"
+          }
+        ]
+      }
+    ]
+  }
+]
+```
+| Key                | Description                                                                                           | Example                                                                                                    |
+|--------------------|-------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------|
+| `answer_sentences` | Textual information or sentences extracted as answers or relevant information related to the citation.| `"Elon Musk's net worth is estimated to be US$241 billion as of August 2024."`                             |
+| `cite_document`    | List of source documents used in the citation. Each document contains:                                |                                                                                                            |
+|                    | - `document`: Text from the source document.                                                          | `"people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3]"`|
+|                    | - `source_id`: Unique identifier for the source document.                                             | `"6874d990-fedc-42bd-b0be-730bcdd59d26"`                                                                   |
+|                    |  - `entity`: List of recognized entities in the document. Each entity contains:                       |                                                                                                            |
+|                    |  - `word`: Recognized word or phrase.                                                                 | `"US$241 billion"`                                                                                         |
+|                    | - `entity_name` Type of the entity (e.g., `MONEY`, `DATE`).                                           | `"MONEY"`                                                                                                  |
+|                    | - `meta`Metadata about the document:                                                                  |  `[]`                                                                                                      |
+
+##### `print(output.hallucination)` 
+`False`
+
+| Key              | Description                                                         | Example |
+|------------------|---------------------------------------------------------------------|---------|
+| `hallucination`  | Indicates if the output contains hallucinated information.          | `false` |
+
+### `print(output.missing)`
+` [] `
+
+| Key           | Description                                 | Example             |
+|---------------|---------------------------------------------|---------------------|
+| `missing`     | List of entities expected but not found.    | `["$100 USD"]` |
+
+
+## Installation
+
+**From PyPI:**
+
+```bash
+pip install rag-citation
+```
+
+**From Source:**
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/your-username/rag-citation.git 
+   cd rag-citation
+   ```
+2. Install the dependencies:
+   ```bash
+   pip install -r requirements.txt 
+   ```
+
+## Configuration
+
+The `Inference` class can be configured with different models and settings:
+
+- **`spacy_model`:** The spaCy model used for named entity recognition (default: `"en_core_web_sm"`). To use different models, pass:
+  - `"sm"` for `en_core_web_sm`
+  - `"md"` for `en_core_web_md`
+  - `"lg"` for `en_core_web_lg`
+  You can download and install spaCy models [here](https://spacy.io/models).
+
+- **`embedding_model`:** The sentence embedding model from the SentenceTransformers library used for semantic similarity (default: `"all-mpnet-base-v2"`). To use different models, pass:
+  - `"sm"` for `avsolatorio/GIST-small-Embedding-v0`
+  - `"md"` for `avsolatorio/GIST-Embedding-v0`
+  - `"lg"` for `avsolatorio/GIST-large-Embedding-v0`
+  Install SentenceTransformers with: `pip install -U sentence-transformers`  
+  You can explore the models on [Hugging Face](https://huggingface.co/avsolatorio/GIST-Embedding-v0).
+
+- **`threshold`:** The similarity threshold value for semantic matching (current default: `0.88`). You can adjust this value as needed.
+
+
+## Contributing
+
+We welcome contributions! Here’s how you can help:
+
+- **Report Bugs:** Submit issues on GitHub.
+- **Suggest Features:**  Open an issue with your ideas.
+- **Code Contributions:** Fork, make changes, and submit a pull request.
+- **Documentation:** Update and enhance our docs.
+
+
+## License
+
+This project is licensed under the [MIT License](LICENSE).
+
+## Acknowledgements
+
+- [SpaCy](https://spacy.io/)
+- [SentenceTransformers](https://www.sbert.net/)
+- [Huggingface/avsolatorio](https://huggingface.co/avsolatorio)

rag_citation-0.0.1/README.md

@@ -0,0 +1,191 @@
+# RAG Citation: Enhancing AI-Generated Content with Automatic Citations (A Non-LLM Approach)
+
+## Project Overview
+
+RAG Citation is an project that combines Retrieval-Augmented Generation (RAG) with automatic citation generation. This tool is designed to enhance the credibility of AI-generated content by providing relevant citations for the information used in generating responses.
+
+## Key Features
+
+- **Non-LLM Approach:** Utilizes efficient algorithms and NLP techniques for citation generation, making it fast and lightweight.
+- **Semantic Search:** Identifies relevant source documents based on meaning and context rather than just keyword matching.
+- **Named Entity Recognition:**  Extracts and returns relevant named entities from LLM-generated answers, such as people, organizations, money and dates.
+- **Flexible Integration:** Can be easily integrated into rag pipeline.
+- **Hallucination (Beta)** This beta feature identifies instances where the LLM-generated answer contains entities like ["DATE", "MONEY", "CARDINAL", "ORDINAL", "QUANTITY", "TIME"], but these entities cannot be found within the context. If such a mismatch occurs, it flags the result as a potential hallucination.
+
+## Quickstart
+* <b>Langchain example</b>: [langchain.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/3.example-langchain.ipynb)
+* <b>Embeddchain example</b>: [embeddchain.ipynb](https://github.com/rahulanand1103/rag-citation/blob/main/docs/examples/2.example-embeddchain.ipynb)
+
+To get started with `rag-citation`, install it using pip:
+
+```bash
+pip install rag-citation
+```
+
+Here's a basic example demonstrating how to use the library:
+
+```python
+from rag_citation import CiteItem, Inference
+import uuid
+
+## Sample context from vectorDB or semantic search
+documents = [
+    "Elon MuskCEO, Tesla$221.6B$439M (0.20%)Real Time Net Worthas of 8/6/24Reflects change since 5 pm ET of prior trading day. 1 in the world todayPhoto by Martin Schoeller for ForbesAbout Elon MuskElon Musk cofounded six companies, including electric car maker Tesla, rocket producer SpaceX and tunneling startup Boring Company.He owns about 12% of Tesla excluding options, but has pledged more than half his shares as collateral for personal loans of up to $3.5 billion.In early 2024, a Delaware judge voided Musk's 2018 deal to receive options equaling an additional 9% of Tesla.",
+    "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3] Musk was born in Pretoria to model Maye and businessman and engineer Errol Musk, and briefly attended the University of Pretoria before immigrating to Canada at age 18, acquiring citizenship through his Canadian-born mother. Two years later, he matriculated at Queen's University at Kingston in Canada. Musk later transferred to the University of Pennsylvania and received bachelor's degrees in economics and physics."
+]
+
+## Example answer generated by an LLM
+answer = "Elon Musk's net worth is estimated to be US$241 billion as of August 2024."
+
+## Helper function to generate a UUID
+def generate_uuid():
+    return str(uuid.uuid4())
+
+## Helper function to create context in the correct format
+def format_document(documents):
+    context = []
+    for document in documents:
+        context.append(
+            {
+                "source_id": generate_uuid(), 
+                "document": document, 
+                "meta": [{"meta-data": "some-info"}], 
+            }
+        )
+    return context
+
+context = format_document(documents)
+cite_item = CiteItem(answer=answer, context=context)
+
+## Initialize the Inference 
+inference = Inference(spacy_model="sm", embedding_model="md")
+
+## Get citation and other information
+output = inference(cite_item)
+
+print("------ Citation ------")
+print(output.citation) 
+print("------ Hallucination ------") 
+print(output.hallucination) 
+print("------ Missing Entities ------")
+print(output.missing) 
+
+```
+
+## Output Explanation
+
+### `print(output.citation)`
+```
+[
+  {
+    "answer_sentences": "Elon Musk's net worth is estimated to be US$241 billion as of August 2024.",
+    "cite_document": [
+      {
+        "document": "people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3]",
+        "source_id": "23d1f1f0-2afa-4749-8639-78ec685fd837",
+        "entity": [
+          {
+            "word": "US$241 billion",
+            "entity_name": "MONEY"
+          },
+          {
+            "word": "August 2024",
+            "entity_name": "DATE"
+          }
+        ],
+        "meta": [
+          {
+            "url": "https://www.forbes.com/profile/elon-musk/",
+            "chunk_id": "1eab8dd1ffa92906f7fc839862871ca5"
+          }
+        ]
+      }
+    ]
+  }
+]
+```
+| Key                | Description                                                                                           | Example                                                                                                    |
+|--------------------|-------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------|
+| `answer_sentences` | Textual information or sentences extracted as answers or relevant information related to the citation.| `"Elon Musk's net worth is estimated to be US$241 billion as of August 2024."`                             |
+| `cite_document`    | List of source documents used in the citation. Each document contains:                                |                                                                                                            |
+|                    | - `document`: Text from the source document.                                                          | `"people in the world; as of August 2024[update], Forbes estimates his net worth to be US$241 billion.[3]"`|
+|                    | - `source_id`: Unique identifier for the source document.                                             | `"6874d990-fedc-42bd-b0be-730bcdd59d26"`                                                                   |
+|                    |  - `entity`: List of recognized entities in the document. Each entity contains:                       |                                                                                                            |
+|                    |  - `word`: Recognized word or phrase.                                                                 | `"US$241 billion"`                                                                                         |
+|                    | - `entity_name` Type of the entity (e.g., `MONEY`, `DATE`).                                           | `"MONEY"`                                                                                                  |
+|                    | - `meta`Metadata about the document:                                                                  |  `[]`                                                                                                      |
+
+##### `print(output.hallucination)` 
+`False`
+
+| Key              | Description                                                         | Example |
+|------------------|---------------------------------------------------------------------|---------|
+| `hallucination`  | Indicates if the output contains hallucinated information.          | `false` |
+
+### `print(output.missing)`
+` [] `
+
+| Key           | Description                                 | Example             |
+|---------------|---------------------------------------------|---------------------|
+| `missing`     | List of entities expected but not found.    | `["$100 USD"]` |
+
+
+## Installation
+
+**From PyPI:**
+
+```bash
+pip install rag-citation
+```
+
+**From Source:**
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/your-username/rag-citation.git 
+   cd rag-citation
+   ```
+2. Install the dependencies:
+   ```bash
+   pip install -r requirements.txt 
+   ```
+
+## Configuration
+
+The `Inference` class can be configured with different models and settings:
+
+- **`spacy_model`:** The spaCy model used for named entity recognition (default: `"en_core_web_sm"`). To use different models, pass:
+  - `"sm"` for `en_core_web_sm`
+  - `"md"` for `en_core_web_md`
+  - `"lg"` for `en_core_web_lg`
+  You can download and install spaCy models [here](https://spacy.io/models).
+
+- **`embedding_model`:** The sentence embedding model from the SentenceTransformers library used for semantic similarity (default: `"all-mpnet-base-v2"`). To use different models, pass:
+  - `"sm"` for `avsolatorio/GIST-small-Embedding-v0`
+  - `"md"` for `avsolatorio/GIST-Embedding-v0`
+  - `"lg"` for `avsolatorio/GIST-large-Embedding-v0`
+  Install SentenceTransformers with: `pip install -U sentence-transformers`  
+  You can explore the models on [Hugging Face](https://huggingface.co/avsolatorio/GIST-Embedding-v0).
+
+- **`threshold`:** The similarity threshold value for semantic matching (current default: `0.88`). You can adjust this value as needed.
+
+
+## Contributing
+
+We welcome contributions! Here’s how you can help:
+
+- **Report Bugs:** Submit issues on GitHub.
+- **Suggest Features:**  Open an issue with your ideas.
+- **Code Contributions:** Fork, make changes, and submit a pull request.
+- **Documentation:** Update and enhance our docs.
+
+
+## License
+
+This project is licensed under the [MIT License](LICENSE).
+
+## Acknowledgements
+
+- [SpaCy](https://spacy.io/)
+- [SentenceTransformers](https://www.sbert.net/)
+- [Huggingface/avsolatorio](https://huggingface.co/avsolatorio)

rag_citation-0.0.1/rag_citation/__init__.py

@@ -0,0 +1,4 @@
+from rag_citation.cite_item import CiteItem
+from rag_citation.inference import Inference
+
+__all__ = ["CiteItem", "Inference"]

rag_citation-0.0.1/rag_citation/base_model/__init__.py

@@ -0,0 +1,4 @@
+from rag_citation.base_model.spacy_model import SpacyBaseModel
+from rag_citation.base_model.embedding_model import EmbedddingBaseModel
+
+__all__ = ["SpacyBaseModel", "EmbedddingBaseModel"]

rag_citation-0.0.1/rag_citation/base_model/embedding_model.py

@@ -0,0 +1,49 @@
+from sentence_transformers import SentenceTransformer
+
+
+class EmbedddingBaseModel:
+    """
+    Base class for embedding models.
+
+    This class provides a shared instance of a SentenceTransformer model based on the specified size.
+
+    Attributes:
+        _model (SentenceTransformer): Static attribute to store the shared SentenceTransformer instance.
+
+    Args:
+        embedding_model (str, optional): Size of the embedding model.
+                                         Choose from "sm" (small), "md" (medium), or "lg" (large).
+                                         Defaults to "sm".
+
+    default model: `sm`
+    """
+
+    _model = None
+
+    def __init__(self, embedding_model="sm"):
+
+        if EmbedddingBaseModel._model is None:
+
+            if embedding_model == "sm":
+                EmbedddingBaseModel._model = SentenceTransformer(
+                    "avsolatorio/GIST-small-Embedding-v0", revision=None
+                )
+
+            elif embedding_model == "md":
+                EmbedddingBaseModel._model = SentenceTransformer(
+                    "avsolatorio/GIST-Embedding-v0", revision=None
+                )
+
+            elif embedding_model == "lg":
+                EmbedddingBaseModel._model = SentenceTransformer(
+                    "avsolatorio/GIST-large-Embedding-v0", revision=None
+                )
+
+            else:
+                print("Warning::please choose `small`, `medium`, or `large`")
+                print("Warning::choosing default model: small")
+                EmbedddingBaseModel._model = SentenceTransformer(
+                    "avsolatorio/GIST-small-Embedding-v0", revision=None
+                )
+
+        self.model = EmbedddingBaseModel._model

rag_citation-0.0.1/rag_citation/base_model/spacy_model.py

@@ -0,0 +1,41 @@
+import spacy
+
+
+class SpacyBaseModel:
+    _nlp = None
+
+    def __init__(self, spacy_model="sm"):
+        """
+        Base class for loading and sharing a SpaCy language model.
+
+        This class ensures that only one instance of the specified SpaCy model is loaded,
+        regardless of how many times the class is instantiated. This can save memory and time.
+
+        Attributes:
+            _nlp (spacy.Language): Static attribute to store the shared SpaCy language model.
+
+        Args:
+            spacy_model (str, optional): Size of the SpaCy model to load.
+                                         Choose from "sm" (small), "md" (medium), or "lg" (large).
+                                         Defaults to "sm".
+
+        default model: `sm`
+        """
+
+        if SpacyBaseModel._nlp is None:
+
+            if spacy_model == "sm":
+                SpacyBaseModel._nlp = spacy.load("en_core_web_sm")
+
+            elif spacy_model == "md":
+                SpacyBaseModel._nlp = spacy.load("en_core_web_md")
+
+            elif spacy_model == "lg":
+                SpacyBaseModel._nlp = spacy.load("en_core_web_lg")
+
+            else:
+                print("Warning::Please provide correct input: sm, md, or lg")
+                print("Warning::Running use sm(en_core_web_sm)")
+                SpacyBaseModel._nlp = spacy.load("en_core_web_sm")
+
+        self.nlp = SpacyBaseModel._nlp

rag_citation-0.0.1/rag_citation/cite_item.py

@@ -0,0 +1,73 @@
+import re
+
+
+class CiteItem:
+    """
+    Data structure to hold an answer, its corresponding context, and metadata.
+
+    Args:
+        answer (str): The answer extracted from the context.
+        context (list): A list of dictionaries, each representing a source of information.
+                        Each dictionary must contain:
+                            - 'document': str, The source text.
+                        Each dictionary may optionally contain:
+                            - 'source_id': str, A unique identifier for the source.
+                            - 'meta': dict, Additional metadata associated with the source.
+
+    Attributes:
+        answer (str): The answer extracted from the context.
+        context (list): The list of context dictionaries.
+        meta (dict): A dictionary mapping source_ids to their corresponding metadata.
+
+    Raises:
+        ValueError: If answer is not a non-empty string, context is not a non-empty list,
+                    or any item in the context list doesn't contain a 'document' key.
+
+    Example:
+        >>> context = [
+        ...     {"document": "This is the first sentence.", "source_id": "doc1", "meta": {"doc_id": "doc_id"}},
+        ...     {"document": "This is the second sentence.", "source_id": "doc2", "meta": {"doc_id": "doc_id"}},
+        ... ]
+        >>> cite_item = CiteItem(answer="some answer", context=context)
+        >>> print(cite_item.meta)
+    """
+
+    def __init__(self, answer: str, context: list):
+        # Validate answer input
+        if not isinstance(answer, str) or not answer.strip():
+            raise ValueError("Answer must be a non-empty string.")
+
+        # Validate context input
+        if not isinstance(context, list) or not context:
+            raise ValueError("Context must be a non-empty list.")
+
+        required_keys = {"document"}
+        for item in context:
+            if not isinstance(item, dict):
+                raise ValueError("Each context item must be a dictionary.")
+            missing_keys = required_keys - set(item.keys())
+            if missing_keys:
+                raise ValueError(
+                    f"Each context item must contain the following keys: {', '.join(missing_keys)}"
+                )
+
+        # Assign attributes
+        self.answer = self._clean_text(answer)
+        self.context = [self._clean_text(item) for item in context]
+        self.meta = {item.get("source_id"): item.get("meta") for item in context}
+
+    def _clean_text(self, text):
+        # Clean the input text by removing newline characters
+        if isinstance(text, str):
+            # text = text.replace("\n", "\n ")
+            text = re.sub(r"^\d+\.\s*", "", text, flags=re.MULTILINE)
+            return text
+        elif isinstance(text, dict):
+            return {k: self._clean_text(v) for k, v in text.items()}
+        else:
+            return text
+
+    def __str__(self):
+        # String representation for easy inspection
+        attributes = ", ".join(f"{key}={getattr(self, key)}" for key in self.__dict__)
+        return f"CiteItem({attributes})"

rag_citation-0.0.1/rag_citation/embedding/__init__.py

@@ -0,0 +1,3 @@
+from rag_citation.embedding.embedding_model import EmbeddingModel
+
+__all__ = ["EmbeddingModel"]

rag_citation-0.0.1/rag_citation/embedding/embedding_model.py

@@ -0,0 +1,33 @@
+from sentence_transformers import SentenceTransformer
+from rag_citation.base_model import EmbedddingBaseModel
+
+
+class EmbeddingModel(EmbedddingBaseModel):
+    """
+    Class for generating sentence embeddings using a shared SentenceTransformer model.
+
+    This class inherits from EmbedddingBaseModel to utilize a shared instance of a
+    SentenceTransformer model. It provides a method for generating embeddings for a
+    given sentence.
+
+    Args:
+        embedding_model (str, optional): Size of the embedding model.
+                                         Choose from "sm" (small), "md" (medium), or "lg" (large).
+                                         Defaults to "sm".
+    """
+
+    def __init__(self, embedding_model="sm"):
+        super().__init__(embedding_model)
+
+    def embedding(self, sentence: str):
+        """
+        Generates embeddings for a given sentence.
+
+        Args:
+            sentence (str): The sentence to embed.
+
+        Returns:
+            torch.Tensor: A tensor representing the sentence embedding.
+        """
+        embeddings = self.model.encode([sentence], convert_to_tensor=True)
+        return embeddings

rag_citation-0.0.1/rag_citation/focus_word/__init__.py

@@ -0,0 +1,4 @@
+from rag_citation.focus_word.focus_word import FocusWord
+from rag_citation.focus_word.schema import FocusWordDataType
+
+__all__ = ["FocusWord", "FocusWordDataType"]