Leveraging Custom Models in RAG Citation: A Path to Enhanced Citation Generation

In the realm of information retrieval and generation, having control over the models you use is critical. As we continue to innovate within the Retrieval-Augmented Generation (RAG) paradigm, integrating custom embedding models into our workflows can significantly enhance the quality and relevance of the generated citations.

What is RAG Citation?

RAG Citation is an advanced tool designed to bridge the gap between AI-generated content and credible sourcing. By augmenting the generation process with relevant citations, we ensure that the information provided is not only accurate but also verifiable. For a deeper dive into RAG Citation and its capabilities, check out my previous post here.

Why Use Custom Models?

1. Enhanced Flexibility: Custom embedding models allow users to tailor the model’s architecture and training to specific datasets or domains. This ensures that the embeddings generated are more aligned with the user’s unique requirements.
2. Improved Accuracy: By utilizing models that have been fine-tuned on domain-specific data, the accuracy of the generated citations improves significantly, which is crucial for maintaining the credibility of generated content.
3. Integration of Domain Knowledge: Custom models can incorporate specialized knowledge that generic models might overlook, further enhancing the contextual relevance of the citations generated.

Implementing a Custom Embedding Model

Integrating a custom embedding model into RAG Citation is straightforward. Below is an example that demonstrates how you can use a custom embedding model for citation generation.

Step 1: Define Your Documents

First, create a list of documents and the answer generated by your language model:

documents = [
    "Elon Musk is the CEO of Tesla and SpaceX...",
    "As of August 2024, Elon Musk's net worth is estimated to be US$241 billion...",
]
answer = "Elon Musk's net worth is estimated to be US$241 billion as of August 2024."

Step 2: Generate Context with Unique Identifiers

Each document must have a unique identifier for proper referencing:

import uuid

def generate_uuid():
    return str(uuid.uuid4())

context = []
for document in documents:
    context.append(
        {
            "source_id": generate_uuid(),
            "document": document,
            "meta": [
                {
                    "url": "https://www.forbes.com/profile/elon-musk/",
                    "chunk_id": generate_uuid(),
                }
            ],
        }
    )

Step 3: Create Your Custom Embedding Model

Now, define your custom embedding model:

from rag_citation.base_model import BaseEmbeddingModel
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, AutoModel
import torch.nn.functional as F
from torch import Tensor

class CustomEmbeddingModel(BaseEmbeddingModel):
    def __init__(self):
        self.tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-small")
        self.model = AutoModel.from_pretrained("intfloat/multilingual-e5-small")

    def average_pool(self, last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
        last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

    def embedding(self, sentence: str):
        query = f"query: {sentence}"
        batch_dict = self.tokenizer(query, max_length=512, padding=True, truncation=True, return_tensors="pt")
        outputs = self.model(**batch_dict)
        embeddings = self.average_pool(outputs.last_hidden_state, batch_dict["attention_mask"])
        embeddings = F.normalize(embeddings, p=2, dim=1)
        return embeddings

Step 4: Utilize RAG Citation with Your Custom Model

Finally, use your custom model to generate citations:

from rag_citation import CiteItem, Inference

cite_item = CiteItem(answer=answer, context=context)
inference = Inference(spacy_model="lg", embedding_model=CustomEmbeddingModel())

output = inference(cite_item)

print("------ Citation --------")
print(output.citation)

print("------ Missing Word --------")
print(output.missing_word)

print("------ Hallucination --------")
print(output.hallucination)

Conclusion

By leveraging custom models within the RAG Citation framework, you can enhance the quality and relevance of your citations. This capability allows for a more personalized and accurate citation generation process, ultimately leading to improved trust in AI-generated content.

Explore more about RAG Citation and start building your custom embedding models today!