45 lines
1.2 KiB
Python
45 lines
1.2 KiB
Python
from modaic import PrecompiledAgent, PrecompiledConfig
|
|
from agent import SignatureGenerator
|
|
import dspy
|
|
|
|
|
|
class PromptToSignatureConfig(PrecompiledConfig):
|
|
lm: str = "gemini/gemini-2.5-pro-preview-03-25"
|
|
max_tokens: int = 4096
|
|
temperature: float = 0.7
|
|
|
|
|
|
class PromptToSignatureAgent(PrecompiledAgent):
|
|
config: PromptToSignatureConfig
|
|
|
|
def __init__(self, config: PromptToSignatureConfig, **kwargs):
|
|
super().__init__(config, **kwargs)
|
|
self.signature_generator = SignatureGenerator()
|
|
|
|
lm = dspy.LM(
|
|
model=config.lm,
|
|
max_tokens=config.max_tokens,
|
|
temperature=config.temperature,
|
|
)
|
|
|
|
self.signature_generator.set_lm(lm)
|
|
|
|
def forward(self, prompt: str, as_dict: bool = False) -> dspy.Prediction:
|
|
# returns dspy.Prediction object or dict
|
|
return (
|
|
self.signature_generator.generate_signature(prompt)
|
|
if as_dict
|
|
else self.signature_generator(prompt)
|
|
)
|
|
|
|
|
|
agent = PromptToSignatureAgent(PromptToSignatureConfig())
|
|
|
|
|
|
def main():
|
|
agent.push_to_hub("fadeleke/prompt-to-signature", with_code=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|