janus-llm 4.3.1__py3-none-any.whl → 4.4.5__py3-none-any.whl

janus/cli/self_eval.py

@@ -0,0 +1,147 @@
+from pathlib import Path
+from typing import Optional
+
+import click
+import typer
+from typing_extensions import Annotated
+
+from janus.cli.constants import REFINERS
+from janus.language.naive.registry import CUSTOM_SPLITTERS
+from janus.utils.enums import LANGUAGES
+
+
+def llm_self_eval(
+    input_dir: Annotated[
+        Path,
+        typer.Option(
+            "--input",
+            "-i",
+            help="The directory containing the source code to be evaluated. "
+            "The files should all be in one flat directory.",
+        ),
+    ],
+    language: Annotated[
+        str,
+        typer.Option(
+            "--language",
+            "-l",
+            help="The language of the source code.",
+            click_type=click.Choice(sorted(LANGUAGES)),
+        ),
+    ],
+    output_dir: Annotated[
+        Path,
+        typer.Option("--output", "-o", help="The directory to store the evaluations in."),
+    ],
+    failure_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--failure-directory",
+            "-f",
+            help="The directory to store failure files during translation",
+        ),
+    ] = None,
+    llm_name: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            "-L",
+            help="The custom name of the model set with 'janus llm add'.",
+        ),
+    ] = "gpt-4o",
+    evaluation_type: Annotated[
+        str,
+        typer.Option(
+            "--evaluation-type",
+            "-e",
+            help="Type of output to evaluate.",
+            click_type=click.Choice(["incose", "comments"]),
+        ),
+    ] = "incose",
+    max_prompts: Annotated[
+        int,
+        typer.Option(
+            "--max-prompts",
+            "-m",
+            help="The maximum number of times to prompt a model on one functional block "
+            "before exiting the application. This is to prevent wasting too much money.",
+        ),
+    ] = 10,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite/--preserve",
+            help="Whether to overwrite existing files in the output directory",
+        ),
+    ] = False,
+    temperature: Annotated[
+        float,
+        typer.Option("--temperature", "-t", help="Sampling temperature.", min=0, max=2),
+    ] = 0.7,
+    collection: Annotated[
+        str,
+        typer.Option(
+            "--collection",
+            "-c",
+            help="If set, will put the translated result into a Chroma DB "
+            "collection with the name provided.",
+        ),
+    ] = None,
+    splitter_type: Annotated[
+        str,
+        typer.Option(
+            "-S",
+            "--splitter",
+            help="Name of custom splitter to use",
+            click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
+        ),
+    ] = "file",
+    refiner_types: Annotated[
+        list[str],
+        typer.Option(
+            "-r",
+            "--refiner",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
+            click_type=click.Choice(list(REFINERS.keys())),
+        ),
+    ] = ["JanusRefiner"],
+    eval_items_per_request: Annotated[
+        int,
+        typer.Option(
+            "--eval-items-per-request",
+            "-rc",
+            help="The maximum number of evaluation items per request",
+        ),
+    ] = None,
+    max_tokens: Annotated[
+        int,
+        typer.Option(
+            "--max-tokens",
+            "-M",
+            help="The maximum number of tokens the model will take in. "
+            "If unspecificed, model's default max will be used.",
+        ),
+    ] = None,
+):
+    from janus.converter.evaluate import InlineCommentEvaluator, RequirementEvaluator
+
+    model_arguments = dict(temperature=temperature)
+    refiner_types = [REFINERS[r] for r in refiner_types]
+    kwargs = dict(
+        eval_items_per_request=eval_items_per_request,
+        model=llm_name,
+        model_arguments=model_arguments,
+        source_language=language,
+        max_prompts=max_prompts,
+        max_tokens=max_tokens,
+        splitter_type=splitter_type,
+        refiner_types=refiner_types,
+    )
+    # Setting parser type here
+    if evaluation_type == "incose":
+        evaluator = RequirementEvaluator(**kwargs)
+    elif evaluation_type == "comments":
+        evaluator = InlineCommentEvaluator(**kwargs)
+
+    evaluator.translate(input_dir, output_dir, failure_dir, overwrite, collection)

janus/cli/translate.py

@@ -0,0 +1,192 @@
+from pathlib import Path
+from typing import Optional
+
+import click
+import typer
+from typing_extensions import Annotated
+
+from janus.cli.constants import REFINERS
+from janus.language.naive.registry import CUSTOM_SPLITTERS
+from janus.utils.enums import LANGUAGES
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+def translate(
+    input_dir: Annotated[
+        Path,
+        typer.Option(
+            "--input",
+            "-i",
+            help="The directory containing the source code to be translated. "
+            "The files should all be in one flat directory.",
+        ),
+    ],
+    source_lang: Annotated[
+        str,
+        typer.Option(
+            "--source-language",
+            "-s",
+            help="The language of the source code.",
+            click_type=click.Choice(sorted(LANGUAGES)),
+        ),
+    ],
+    output_dir: Annotated[
+        Path,
+        typer.Option(
+            "--output", "-o", help="The directory to store the translated code in."
+        ),
+    ],
+    target_lang: Annotated[
+        str,
+        typer.Option(
+            "--target-language",
+            "-t",
+            help="The desired output language to translate the source code to. The "
+            "format can follow a 'language-version' syntax.  Use 'text' to get plaintext"
+            "results as returned by the LLM. Examples: `python-3.10`, `mumps`, `java-10`,"
+            "text.",
+        ),
+    ],
+    llm_name: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            "-L",
+            help="The custom name of the model set with 'janus llm add'.",
+        ),
+    ],
+    failure_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--failure-directory",
+            "-f",
+            help="The directory to store failure files during translation",
+        ),
+    ] = None,
+    max_prompts: Annotated[
+        int,
+        typer.Option(
+            "--max-prompts",
+            "-m",
+            help="The maximum number of times to prompt a model on one functional block "
+            "before exiting the application. This is to prevent wasting too much money.",
+        ),
+    ] = 10,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite/--preserve",
+            help="Whether to overwrite existing files in the output directory",
+        ),
+    ] = False,
+    skip_context: Annotated[
+        bool,
+        typer.Option(
+            "--skip-context",
+            help="Prompts will include any context information associated with source"
+            " code blocks, unless this option is specified",
+        ),
+    ] = False,
+    temp: Annotated[
+        float,
+        typer.Option("--temperature", "-T", help="Sampling temperature.", min=0, max=2),
+    ] = 0.7,
+    prompt_template: Annotated[
+        str,
+        typer.Option(
+            "--prompt-template",
+            "-p",
+            help="Name of the Janus prompt template directory or "
+            "path to a directory containing those template files.",
+        ),
+    ] = "simple",
+    collection: Annotated[
+        str,
+        typer.Option(
+            "--collection",
+            "-c",
+            help="If set, will put the translated result into a Chroma DB "
+            "collection with the name provided.",
+        ),
+    ] = None,
+    splitter_type: Annotated[
+        str,
+        typer.Option(
+            "-S",
+            "--splitter",
+            help="Name of custom splitter to use",
+            click_type=click.Choice(list(CUSTOM_SPLITTERS.keys())),
+        ),
+    ] = "file",
+    refiner_types: Annotated[
+        list[str],
+        typer.Option(
+            "-r",
+            "--refiner",
+            help="List of refiner types to use. Add -r for each refiner to use in\
+                refinement chain",
+            click_type=click.Choice(list(REFINERS.keys())),
+        ),
+    ] = ["JanusRefiner"],
+    retriever_type: Annotated[
+        str,
+        typer.Option(
+            "-R",
+            "--retriever",
+            help="Name of custom retriever to use",
+            click_type=click.Choice(["active_usings", "language_docs"]),
+        ),
+    ] = None,
+    max_tokens: Annotated[
+        int,
+        typer.Option(
+            "--max-tokens",
+            "-M",
+            help="The maximum number of tokens the model will take in. "
+            "If unspecificed, model's default max will be used.",
+        ),
+    ] = None,
+    use_janus_inputs: Annotated[
+        bool,
+        typer.Option(
+            "-j",
+            "--use-janus-inputs",
+            help="Prsent if translator should use janus files as inputs",
+        ),
+    ] = False,
+):
+    from janus.cli.constants import db_loc, get_collections_config
+    from janus.converter.translate import Translator
+
+    refiner_types = [REFINERS[r] for r in refiner_types]
+    try:
+        target_language, target_version = target_lang.split("-")
+    except ValueError:
+        target_language = target_lang
+        target_version = None
+    # make sure not overwriting input
+    if source_lang.lower() == target_language.lower() and input_dir == output_dir:
+        log.error("Output files would overwrite input! Aborting...")
+        raise ValueError
+
+    model_arguments = dict(temperature=temp)
+    collections_config = get_collections_config()
+    translator = Translator(
+        model=llm_name,
+        model_arguments=model_arguments,
+        source_language=source_lang,
+        target_language=target_language,
+        target_version=target_version,
+        max_prompts=max_prompts,
+        max_tokens=max_tokens,
+        prompt_templates=prompt_template,
+        db_path=db_loc,
+        db_config=collections_config,
+        splitter_type=splitter_type,
+        refiner_types=refiner_types,
+        retriever_type=retriever_type,
+        use_janus_inputs=use_janus_inputs,
+    )
+    translator.translate(input_dir, output_dir, failure_dir, overwrite, collection)

janus/converter/__init__.py

@@ -1,6 +1,6 @@
 from janus.converter.converter import Converter
 from janus.converter.diagram import DiagramGenerator
-from janus.converter.document import Documenter, MadLibsDocumenter, MultiDocumenter
+from janus.converter.document import ClozeDocumenter, Documenter, MultiDocumenter
 from janus.converter.evaluate import Evaluator
 from janus.converter.partition import Partitioner
 from janus.converter.requirements import RequirementsDocumenter

janus/converter/_tests/test_translate.py

@@ -11,6 +11,7 @@ from janus.converter.diagram import DiagramGenerator
 from janus.converter.requirements import RequirementsDocumenter
 from janus.converter.translate import Translator
 from janus.language.block import CodeBlock, TranslatedCodeBlock
+from janus.refiners.format import CodeFormatRefiner
 
 
 class MockCollection(VectorStore):
@@ -50,6 +51,7 @@ class TestTranslator(unittest.TestCase):
             target_language="python",
             target_version="3.10",
             splitter_type="ast-flex",
+            refiner_types=[CodeFormatRefiner],
         )
         self.test_file = Path("janus/language/treesitter/_tests/languages/fortran.f90")
         self.TEST_FILE_EMBEDDING_COUNT = 14
@@ -57,14 +59,14 @@ class TestTranslator(unittest.TestCase):
         self.req_translator = RequirementsDocumenter(
             model="gpt-4o-mini",
             source_language="fortran",
-            prompt_template="requirements",
+            prompt_templates="requirements",
         )
 
     @pytest.mark.translate
     def test_translate(self):
         """Test translate method."""
         # Delete a file if it's already there
-        python_file = self.test_file.parent / "python" / f"{self.test_file.stem}.py"
+        python_file = self.test_file.parent / "python" / f"{self.test_file.stem}.json"
         python_file.unlink(missing_ok=True)
         python_file.parent.rmdir() if python_file.parent.is_dir() else None
         self.translator.translate(self.test_file.parent, self.test_file.parent / "python")
@@ -80,7 +82,7 @@ class TestTranslator(unittest.TestCase):
         self.assertRaises(
             ValueError, self.translator.set_source_language, "scribbledy-doop"
         )
-        self.translator.set_prompt("pish posh")
+        self.translator.set_prompts(["pish posh"])
         self.assertRaises(ValueError, self.translator._load_parameters)
 
 
@@ -147,10 +149,10 @@ def test_language_combinations(
     translator.set_model("gpt-4o")
     translator.set_source_language(source_language)
     translator.set_target_language(expected_target_language, expected_target_version)
-    translator.set_prompt(prompt_template)
+    translator.set_prompts(prompt_template)
     translator._load_parameters()
     assert translator._target_language == expected_target_language  # nosec
     assert translator._target_version == expected_target_version  # nosec
     assert translator._splitter.language == source_language  # nosec
     assert translator._splitter.model.model_name == "gpt-4o"  # nosec
-    assert translator._prompt_template_name == prompt_template  # nosec
+    assert translator._prompt_template_names == [prompt_template]  # nosec

janus/converter/chain.py

@@ -0,0 +1,180 @@
+from pathlib import Path
+
+from janus.converter.converter import Converter
+from janus.language.block import CodeBlock, TranslatedCodeBlock
+from janus.utils.logger import create_logger
+
+log = create_logger(__name__)
+
+
+class ConverterChain(Converter):
+    """
+    Class for representing multiple converters chained together
+    """
+
+    def __init__(self, *args, **kwargs) -> None:
+        if len(args) == 0:
+            raise ValueError("Error: Converter chain must be passed at least 1 converter")
+        for converter in args:
+            if not isinstance(converter, Converter):
+                raise ValueError(f"Error: unrecognized type: {type(converter)}")
+        self._converters = args
+        kwargs.update(
+            source_language=self._converters[0].source_language,
+            target_language=self._converters[-1]._target_language,
+            target_version=self._converters[-1]._target_version,
+            use_janus_inputs=self._converters[0]._use_janus_inputs,
+        )
+        super().__init__(**kwargs)
+
+    def _run_converters(
+        self, translated_code_block, name: str, failure_path: Path | None = None
+    ):
+        for i, converter in enumerate(self._converters[1:]):
+            if not translated_code_block.translated:
+                log.info(
+                    f"Error: chain failed to translate at step {i}:"
+                    f"{self._converters[i].__class__.__name__}"
+                )
+                break
+            if converter._use_janus_inputs:
+                janus_obj = self._converters[i]._get_output_obj(translated_code_block)
+                translated_code_block = converter.translate_janus_obj(
+                    janus_obj, name, failure_path
+                )
+            else:
+                translated_code_block = converter.translate_block(
+                    translated_code_block.to_codeblock(), name, failure_path
+                )
+        if not translated_code_block.translated:
+            log.info(
+                f"Error: chain failed to translate at step {len(self._converters)-1}: "
+                f"{self._converters[-1].__class__.__name__}"
+            )
+
+        return translated_code_block
+
+    def translate_file(
+        self, file: Path, failure_path: Path | None = None
+    ) -> TranslatedCodeBlock:
+        """Translate a file using the chain of converters
+
+        Arguments:
+            file: The file to translate
+            failure_path: The path to write the failure file to
+
+        Returns:
+            The translated code block
+        """
+        filename = file.name
+        translated_code_block = self._converters[0].translate_file(file, failure_path)
+        translated_code_block = self._run_converters(
+            translated_code_block, filename, failure_path
+        )
+        return translated_code_block
+
+    def translate_text(
+        self, text: str, name: str, failure_path: Path | None = None
+    ) -> TranslatedCodeBlock:
+        """Translate a text using the chain of converters
+
+        Arguments:
+            text: The text to translate
+            name: The name of the file
+            failure_path: The path to write the failure file to
+
+        Returns:
+            The translated code block
+        """
+        translated_code_block = self._converters[0].translate_text(
+            text, name, failure_path
+        )
+        translated_code_block = self._run_converters(
+            translated_code_block, name, failure_path
+        )
+        return translated_code_block
+
+    def translate_block(
+        self,
+        input_block: CodeBlock | list[CodeBlock],
+        name: str,
+        failure_path: Path | None = None,
+    ) -> TranslatedCodeBlock:
+        """Translate a block of code using the chain of converters
+
+        Arguments:
+            input_block: The block of code to translate
+            name: The name of the file
+            failure_path: The path to write the failure file to
+
+        Returns:
+            The translated code block
+        """
+        translated_code_block = self._converters[0].translate_block(
+            input_block, name, failure_path
+        )
+        translated_code_block = self._run_converters(
+            translated_code_block, name, failure_path
+        )
+        return translated_code_block
+
+    def _get_output_obj(
+        self, block: TranslatedCodeBlock | list, combine_children: bool = True
+    ) -> dict[str, int | float | str | dict[str, str] | dict[str, float]]:
+        output_obj = super()._get_output_obj(block, combine_children)
+        intermediate_outputs = []
+        for i, intermediate_out in enumerate(block.previous_generations):
+            if isinstance(intermediate_out, TranslatedCodeBlock):
+                intermediate_outputs.append(
+                    self._converters[i]._get_output_obj(intermediate_out)
+                )
+            else:
+                intermediate_outputs.append(intermediate_out)
+        intermediate_outputs.append(self._converters[-1]._get_output_obj(block))
+        output_obj["intermediate_outputs"] = intermediate_outputs
+        metadata = output_obj["metadata"]
+        metadata["cost"] += sum(
+            b.cost if isinstance(b, TranslatedCodeBlock) else b["metadata"]["cost"]
+            for b in block.previous_generations
+        )
+        metadata["processing_time"] += sum(
+            (
+                b.processing_time
+                if isinstance(b, TranslatedCodeBlock)
+                else b["metadata"]["processing_time"]
+            )
+            for b in block.previous_generations
+        )
+        metadata["num_requests"] += sum(
+            (
+                b.total_num_requests
+                if isinstance(b, TranslatedCodeBlock)
+                else b["metadata"]["num_requests"]
+            )
+            for b in block.previous_generations
+        )
+        metadata["input_tokens"] += sum(
+            (
+                b.total_request_input_tokens
+                if isinstance(b, TranslatedCodeBlock)
+                else b["metadata"]["input_tokens"]
+            )
+            for b in block.previous_generations
+        )
+        metadata["output_tokens"] += sum(
+            (
+                b.total_request_output_tokens
+                if isinstance(b, TranslatedCodeBlock)
+                else b["metadata"]["output_tokens"]
+            )
+            for b in block.previous_generations
+        )
+        output_obj["metadata"] = metadata
+        if len(block.previous_generations) > 0:
+            b = block.previous_generations[0]
+            output_obj["input"] = (
+                (b.original.text or "")
+                if isinstance(b, TranslatedCodeBlock)
+                else b["input"]
+            )
+        return output_obj