36 lines
945 B
Python
36 lines
945 B
Python
from modaic import PrecompiledAgent, PrecompiledConfig
|
|
import dspy
|
|
|
|
|
|
class ReceiptClassifierConfig(
|
|
PrecompiledConfig
|
|
): # Configurable for each agent instance
|
|
lm: str = "openai/gpt-4o-mini" # Set OPENAI_API_KEY
|
|
max_tokens: int = 500
|
|
|
|
|
|
class ReceiptClassifier(PrecompiledAgent): # Extends dspy.module()
|
|
config: ReceiptClassifierConfig
|
|
|
|
def __init__(self, config: ReceiptClassifierConfig, **kwargs):
|
|
super().__init__(config, **kwargs)
|
|
|
|
self.extract = dspy.Predict("image_url -> total_cost")
|
|
lm = dspy.LM(config.lm, max_tokens=config.max_tokens)
|
|
self.extract.set_lm(lm)
|
|
|
|
def forward(self, image_url: str) -> str:
|
|
receipt_image = dspy.Image.from_url(image_url)
|
|
return self.extract(receipt_image)
|
|
|
|
|
|
agent = ReceiptClassifier(ReceiptClassifierConfig())
|
|
|
|
|
|
def main():
|
|
agent.push_to_hub("farouk1/receipt-classifier", with_code=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|