35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
import dspy
|
|
from modaic import PrecompiledAgent, PrecompiledConfig
|
|
|
|
|
|
# DSPy code
|
|
class GenerateAnswer(dspy.Signature):
|
|
"""Generate answers to the questions"""
|
|
|
|
context = dspy.InputField(
|
|
desc="may contain relevant facts and psychological insights"
|
|
)
|
|
question = dspy.InputField()
|
|
answer_choices = dspy.InputField()
|
|
answer = dspy.OutputField(desc="often between 1 and 5 words")
|
|
|
|
|
|
class CoTSimplifiedBaleenConfig(PrecompiledConfig):
|
|
model: str = "gpt-3.5-turbo"
|
|
max_tokens: int = 1000
|
|
|
|
|
|
class CoTSimplifiedBaleen(PrecompiledAgent):
|
|
config: CoTSimplifiedBaleenConfig
|
|
|
|
def __init__(self, config: CoTSimplifiedBaleenConfig, **kwargs):
|
|
super().__init__(config, **kwargs)
|
|
self.generate_answer = dspy.ChainOfThought(GenerateAnswer)
|
|
self.generate_answer.set_lm(dspy.LM(model=config.model, max_tokens=config.max_tokens))
|
|
|
|
def forward(self, question, context, answer_choices):
|
|
pred = self.generate_answer(
|
|
context=context, question=question, answer_choices=answer_choices
|
|
)
|
|
return dspy.Prediction(context=context, answer=pred.answer)
|