fix nn4k error

This commit is contained in:
Qu 2024-01-07 23:36:28 +08:00
parent 03d19c4053
commit cebca72c43
2 changed files with 2 additions and 2 deletions

View File

@ -66,7 +66,7 @@ class LLMBasedExtractor(SPGExtractor):
def to_rest(self):
"""Transforms `LLMBasedExtractor` to REST model `ExtractNodeConfig`."""
params = dict()
params["model_config"] = json.dumps(self.llm._nn_config)
params["model_config"] = json.dumps(self.llm.init_args)
params["prompt_config"] = json.dumps(
[OperatorClient().serialize(op.to_rest()) for op in self.prompt_ops]
)

View File

@ -165,7 +165,7 @@ class PromptOp(BaseOp, ABC):
) -> List[Dict[str, str]]:
if isinstance(response, list) and len(response) > 0:
response = response[0]
variables.update({f"{self.name}": response})
variables.update({f"{self.__class__.__name__}": response})
return [variables]