37 lines
1.3 KiB
Python
37 lines
1.3 KiB
Python
import dspy
|
|
from modaic import PrecompiledProgram, PrecompiledConfig
|
|
|
|
class CodeGeneratorConfig(PrecompiledConfig):
|
|
model : str = "openai//models/checkpoint-5"
|
|
api_base : str = "https://modaic-ai--grpo-demo-serve.modal.run/v1"
|
|
max_tokens : int = 10000
|
|
temperature : float = 0.7
|
|
|
|
|
|
class CodeGeneration(dspy.Signature):
|
|
query: str = dspy.InputField(desc="The query to generate code for.")
|
|
code: str = dspy.OutputField(desc="The code to generate as a python function.")
|
|
|
|
class CodeGenerator(PrecompiledProgram):
|
|
config: CodeGeneratorConfig
|
|
|
|
def __init__(self, config: CodeGeneratorConfig, **kwargs):
|
|
super().__init__(config=config, **kwargs)
|
|
|
|
modal_lm = dspy.LM(
|
|
model=config.model,
|
|
api_base=config.api_base,
|
|
max_tokens=config.max_tokens,
|
|
temperature=config.temperature,
|
|
)
|
|
self.answer_question = dspy.Predict(CodeGeneration)
|
|
self.answer_question.set_lm(modal_lm)
|
|
|
|
def forward(self, query):
|
|
return self.answer_question(query=query)
|
|
|
|
code_generator = CodeGenerator(CodeGeneratorConfig())
|
|
print(code_generator(query="Write a python function that returns the sum of two numbers.").code)
|
|
code_generator.push_to_hub("modaic/code-generator-trl-grpo", with_code=True, tag="v2.0.2", commit_message="Add README.md")
|
|
|
|
|