RAG Dataset Annotation Pipeline

annotation/annotate.py

@@ -0,0 +1,61 @@
+import time
+from typing import Annotated
+from pydantic import BaseModel, Field
+
+from langgraph.graph import StateGraph, START, END
+from langchain_model import ChatModel
+
+from prompts import ANNOTATION_PROMPT, ANNOTATION_PARSER
+
+
+# -------------------------------------------------------------------
+# States
+# -------------------------------------------------------------------
+class InputState(BaseModel):
+    context: str = Field(default="")
+    claim: str = Field(default="")
+
+
+class AnnotationState(BaseModel):
+    context: str = Field(default="")
+    claim: str = Field(default="")
+    faithfulness: str = Field(default="")
+    factuality: str = Field(default="")
+
+
+class OutputState(BaseModel):
+    faithfulness: str = Field(default="")
+    factuality: str = Field(default="")
+
+
+# -------------------------------------------------------------------
+# Model Initialization
+# -------------------------------------------------------------------
+model = ChatModel()
+
+# -------------------------------------------------------------------
+# Chain Initialization
+# -------------------------------------------------------------------
+chain = ANNOTATION_PROMPT | model | ANNOTATION_PARSER
+
+
+# -------------------------------------------------------------------
+# Annotation Functions
+# -------------------------------------------------------------------
+def annotate(state: AnnotationState) -> AnnotationState:
+    result = chain.invoke({"context": state.context, "claim": state.claim})
+    state.faithfulness = result.faithfulness
+    state.factuality = result.factuality
+    return state
+
+
+# -------------------------------------------------------------------
+# Main Graph
+# -------------------------------------------------------------------
+builder = StateGraph(AnnotationState, input=InputState, output=OutputState)
+builder.add_node("annotate", annotate)
+builder.add_edge(START, "annotate")
+builder.add_edge("annotate", END)
+
+graph = builder.compile()
+graph.step_timeout = 20