30 lines
883 B
Python
30 lines
883 B
Python
# Here you can define your dspy.Module as a modaic.PrecompiledProgram
|
|
from typing import Literal
|
|
|
|
import dspy
|
|
from modaic import PrecompiledProgram
|
|
|
|
DocumentLabel = Literal["legal", "healthcare", "personal", "tax", "other"]
|
|
|
|
|
|
class ExtractSignature(dspy.Signature):
|
|
"""Classify a document snippet into one of the supported label buckets."""
|
|
|
|
context: str = dspy.InputField(desc="Document text or snippet to classify")
|
|
label: DocumentLabel = dspy.OutputField(
|
|
desc="Return exactly one of: legal, healthcare, personal, tax, or other"
|
|
)
|
|
|
|
|
|
class DocExtract(PrecompiledProgram):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.cot = dspy.ChainOfThought(ExtractSignature)
|
|
|
|
def forward(self, context: str) -> dspy.Prediction:
|
|
return self.cot(context=context)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
DocExtract().push_to_hub("tyrin/doc-extract")
|