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

janus/cli/embedding.py

@@ -0,0 +1,122 @@
+import click
+import typer
+from typing_extensions import Annotated
+
+from janus.embedding.embedding_models_info import EmbeddingModelType
+
+embedding = typer.Typer(
+    help="Embedding model commands",
+    add_completion=False,
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+
+@embedding.command("add", help="Add an embedding model config to janus")
+def embedding_add(
+    model_name: Annotated[
+        str, typer.Argument(help="The user's custom name for the model")
+    ],
+    model_type: Annotated[
+        str,
+        typer.Option(
+            "--type",
+            "-t",
+            help="The type of the model",
+            click_type=click.Choice(list(val.value for val in EmbeddingModelType)),
+        ),
+    ] = "OpenAI",
+):
+    import json
+    from pathlib import Path
+
+    from pydantic import AnyHttpUrl
+
+    from janus.embedding.embedding_models_info import (
+        EMBEDDING_COST_PER_MODEL,
+        EMBEDDING_MODEL_CONFIG_DIR,
+        EMBEDDING_TOKEN_LIMITS,
+    )
+
+    if not EMBEDDING_MODEL_CONFIG_DIR.exists():
+        EMBEDDING_MODEL_CONFIG_DIR.mkdir(parents=True)
+    model_cfg = EMBEDDING_MODEL_CONFIG_DIR / f"{model_name}.json"
+    if model_type in EmbeddingModelType.HuggingFaceInferenceAPI.values:
+        hf = typer.style("HuggingFaceInferenceAPI", fg="yellow")
+        url = typer.prompt(f"Enter the {hf} model's URL", type=str, value_proc=AnyHttpUrl)
+        api_model_name = typer.prompt("Enter the model's name", type=str, default="")
+        api_key = typer.prompt("Enter the API key", type=str, default="")
+        max_tokens = typer.prompt(
+            "Enter the model's maximum tokens", default=8191, type=int
+        )
+        in_cost = typer.prompt("Enter the cost per input token", default=0, type=float)
+        out_cost = typer.prompt("Enter the cost per output token", default=0, type=float)
+        params = dict(
+            model_name=api_model_name,
+            api_key=api_key,
+        )
+        cfg = {
+            "model_type": model_type,
+            "model_identifier": str(url),
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": {"input": in_cost, "output": out_cost},
+        }
+    elif model_type in EmbeddingModelType.HuggingFaceLocal.values:
+        hf = typer.style("HuggingFace", fg="yellow")
+        model_id = typer.prompt(
+            f"Enter the {hf} model ID",
+            default="sentence-transformers/all-MiniLM-L6-v2",
+            type=str,
+        )
+        cache_folder = str(
+            Path(
+                typer.prompt(
+                    "Enter the model's cache folder",
+                    default=EMBEDDING_MODEL_CONFIG_DIR / "cache",
+                    type=str,
+                )
+            )
+        )
+        max_tokens = typer.prompt(
+            "Enter the model's maximum tokens", default=8191, type=int
+        )
+        params = dict(
+            cache_folder=str(cache_folder),
+        )
+        cfg = {
+            "model_type": model_type,
+            "model_identifier": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": {"input": 0, "output": 0},
+        }
+    elif model_type in EmbeddingModelType.OpenAI.values:
+        available_models = list(EMBEDDING_COST_PER_MODEL.keys())
+
+        open_ai = typer.style("OpenAI", fg="green")
+        prompt = f"Enter the {open_ai} model name"
+
+        model_name = typer.prompt(
+            prompt,
+            default="text-embedding-3-small",
+            type=click.types.Choice(available_models),
+            show_choices=False,
+        )
+        params = dict(
+            model=model_name,
+        )
+        max_tokens = EMBEDDING_TOKEN_LIMITS[model_name]
+        model_cost = EMBEDDING_COST_PER_MODEL[model_name]
+        cfg = {
+            "model_type": model_type,
+            "model_identifier": model_name,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+        }
+    else:
+        raise ValueError(f"Unknown model type {model_type}")
+    with open(model_cfg, "w") as f:
+        json.dump(cfg, f, indent=2)
+    print(f"Model config written to {model_cfg}")

janus/cli/llm.py

@@ -0,0 +1,187 @@
+import click
+import typer
+from typing_extensions import Annotated
+
+from janus.llm.models_info import MODEL_TYPE_CONSTRUCTORS
+
+llm = typer.Typer(
+    help="LLM commands",
+    add_completion=False,
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+
+@llm.command("add", help="Add a model config to janus")
+def llm_add(
+    model_name: Annotated[
+        str, typer.Argument(help="The user's custom name of the model")
+    ],
+    model_type: Annotated[
+        str,
+        typer.Option(
+            "--type",
+            "-t",
+            help="The type of the model",
+            click_type=click.Choice(sorted(list(MODEL_TYPE_CONSTRUCTORS.keys()))),
+        ),
+    ] = "Azure",
+):
+    import json
+
+    from janus.llm.models_info import (
+        COST_PER_1K_TOKENS,
+        MODEL_CONFIG_DIR,
+        MODEL_ID_TO_LONG_ID,
+        TOKEN_LIMITS,
+        azure_models,
+        bedrock_models,
+        openai_models,
+    )
+
+    if not MODEL_CONFIG_DIR.exists():
+        MODEL_CONFIG_DIR.mkdir(parents=True)
+    model_cfg = MODEL_CONFIG_DIR / f"{model_name}.json"
+    if model_type == "HuggingFace":
+        url = typer.prompt("Enter the model's URL")
+        max_tokens = typer.prompt(
+            "Enter the model's maximum tokens", default=4096, type=int
+        )
+        in_cost = typer.prompt("Enter the cost per input token", default=0, type=float)
+        out_cost = typer.prompt("Enter the cost per output token", default=0, type=float)
+        params = dict(
+            inference_server_url=url,
+            max_new_tokens=max_tokens,
+            top_k=10,
+            top_p=0.95,
+            typical_p=0.95,
+            temperature=0.01,
+            repetition_penalty=1.03,
+            timeout=240,
+        )
+        cfg = {
+            "model_type": model_type,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": {"input": in_cost, "output": out_cost},
+            "input_token_proportion": 0.4,
+        }
+    elif model_type == "HuggingFaceLocal":
+        model_id = typer.prompt("Enter the model ID")
+        task = typer.prompt("Enter the task")
+        max_tokens = typer.prompt(
+            "Enter the model's maximum tokens", default=4096, type=int
+        )
+        in_cost = 0
+        out_cost = 0
+        params = {"model_id": model_id, "task": task}
+        cfg = {
+            "model_type": model_type,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": {"input": in_cost, "output": out_cost},
+            "input_token_proportion": 0.4,
+        }
+    elif model_type == "OpenAI":
+        print("DEPRECATED: Use 'Azure' instead. CTRL+C to exit.")
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="gpt-4o",
+            type=click.Choice(openai_models),
+            show_choices=False,
+        )
+        params = dict(
+            model_name=model_name,
+            temperature=0.7,
+            n=1,
+        )
+        max_tokens = TOKEN_LIMITS[model_name]
+        model_cost = COST_PER_1K_TOKENS[model_name]
+        cfg = {
+            "model_type": model_type,
+            "model_id": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+            "input_token_proportion": 0.4,
+        }
+    elif model_type == "Azure":
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="gpt-4o",
+            type=click.Choice(azure_models),
+            show_choices=False,
+        )
+        params = dict(
+            # Azure uses the "azure_deployment" key for what we're calling "long_model_id"
+            azure_deployment=MODEL_ID_TO_LONG_ID[model_id],
+            temperature=0.7,
+            n=1,
+        )
+        max_tokens = TOKEN_LIMITS[MODEL_ID_TO_LONG_ID[model_id]]
+        model_cost = COST_PER_1K_TOKENS[MODEL_ID_TO_LONG_ID[model_id]]
+        cfg = {
+            "model_type": model_type,
+            "model_id": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+            "input_token_proportion": 0.4,
+        }
+    elif model_type == "BedrockChat" or model_type == "Bedrock":
+        model_id = typer.prompt(
+            "Enter the model ID (list model IDs with `janus llm ls -a`)",
+            default="bedrock-claude-sonnet",
+            type=click.Choice(bedrock_models),
+            show_choices=False,
+        )
+        params = dict(
+            # Bedrock uses the "model_id" key for what we're calling "long_model_id"
+            model_id=MODEL_ID_TO_LONG_ID[model_id],
+            model_kwargs={"temperature": 0.7},
+        )
+        max_tokens = TOKEN_LIMITS[MODEL_ID_TO_LONG_ID[model_id]]
+        model_cost = COST_PER_1K_TOKENS[MODEL_ID_TO_LONG_ID[model_id]]
+        cfg = {
+            "model_type": model_type,
+            "model_id": model_id,
+            "model_args": params,
+            "token_limit": max_tokens,
+            "model_cost": model_cost,
+            "input_token_proportion": 0.4,
+        }
+    else:
+        raise ValueError(f"Unknown model type {model_type}")
+    with open(model_cfg, "w") as f:
+        json.dump(cfg, f, indent=2)
+    print(f"Model config written to {model_cfg}")
+
+
+@llm.command("ls", help="List all of the user-configured models")
+def llm_ls(
+    all: Annotated[
+        bool,
+        typer.Option(
+            "--all",
+            "-a",
+            is_flag=True,
+            help="List all models, including the default model IDs.",
+            click_type=click.Choice(sorted(list(MODEL_TYPE_CONSTRUCTORS.keys()))),
+        ),
+    ] = False,
+):
+    import json
+
+    from janus.cli.constants import MODEL_CONFIG_DIR
+    from janus.llm.models_info import MODEL_TYPES
+
+    print("\n[green]User-configured models[/green]:")
+    for model_cfg in MODEL_CONFIG_DIR.glob("*.json"):
+        with open(model_cfg, "r") as f:
+            cfg = json.load(f)
+        print(f"\t[blue]{model_cfg.stem}[/blue]: [purple]{cfg['model_type']}[/purple]")
+
+    if all:
+        print("\n[green]Available model IDs[/green]:")
+        for model_id, model_type in MODEL_TYPES.items():
+            print(f"\t[blue]{model_id}[/blue]: [purple]{model_type}[/purple]")

janus/cli/partition.py

@@ -0,0 +1,125 @@
+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 partition(
+    input_dir: Annotated[
+        Path,
+        typer.Option(
+            "--input",
+            "-i",
+            help="The directory containing the source code to be partitioned. ",
+        ),
+    ],
+    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-dir", "-o", help="The directory to store the partitioned code in."
+        ),
+    ],
+    llm_name: Annotated[
+        str,
+        typer.Option(
+            "--llm",
+            "-L",
+            help="The custom name of the model set with 'janus llm add'.",
+        ),
+    ] = "gpt-4o",
+    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,
+    temperature: Annotated[
+        float,
+        typer.Option("--temperature", "-t", help="Sampling temperature.", min=0, max=2),
+    ] = 0.7,
+    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"],
+    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,
+    partition_token_limit: Annotated[
+        int,
+        typer.Option(
+            "--partition-tokens",
+            "-pt",
+            help="The limit on the number of tokens per partition.",
+        ),
+    ] = 8192,
+):
+    from janus.converter.partition import Partitioner
+
+    refiner_types = [REFINERS[r] for r in refiner_types]
+    model_arguments = dict(temperature=temperature)
+    kwargs = dict(
+        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,
+        partition_token_limit=partition_token_limit,
+    )
+    partitioner = Partitioner(**kwargs)
+    partitioner.translate(input_dir, output_dir, failure_dir, overwrite)

janus/cli/self_eval.py

@@ -0,0 +1,149 @@
+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-dir", "-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)