[Example] RA-DIT Implement retrieval-augmented instruction finetuning trainer

examples/ra-dit/ra_dit/trainers_and_testers/generator.py

@@ -9,12 +9,51 @@
 
 from fed_rag.decorators import federate
 from fed_rag.types import TestResult, TrainResult
+from fed_rag.types.rag_system import RAGSystem
 
 # Dataset
 train_dataset = load_dataset("stanfordnlp/imdb", split="train[:20]")
 val_dataset = load_dataset("stanfordnlp/imdb", split="test[:10]")
 
 
+# Custom SFTTrainer to implement RA-IT
+class RetrievalAugmentedSFTTrainer(SFTTrainer):
+    """A custom SFTTrainer to implement retrieval-augmented instruction fine-tuning."""
+
+    def compute_loss(self, model: PreTrainedModel, inputs, return_outputs=False):
+        try:
+            rag_system: RAGSystem = model.__associated_rag_system
+        except AttributeError:
+            raise ValueError(
+                "Unable to get associated RAGSystem with supplied `BaseGenerator`."
+            )
+        questions = inputs.pop("questions")
+        answers = inputs.pop("answers")
+
+        # retrieve
+        for q in questions:
+            source_nodes = rag_system.retrieve(query=q)
+            # parallel in-context retrieval augmentation
+            proba_tensors = []
+            for n in source_nodes:
+                context = n.node.text_content
+                # prepare prompt
+                prompt = rag_system.generator.prompt_template.format(
+                    question=q, context=context
+                )
+                # probas
+                outputs = model(prompt)
+                logits = outputs.logits
+                probas = ...
+                proba_tensors.append(probas)
+            weighted_probas = ...
+
+            # loss
+            loss = ...  # masked language model
+
+        return (loss, outputs) if return_outputs else loss
+
+
 @federate.trainer.huggingface
 def generator_train_loop(
     model: PreTrainedModel,