janus-llm 4.3.1__py3-none-any.whl → 4.4.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,191 @@
+import click
+import typer
+from rich import print
+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 token limit", default=65536, type=int
+        )
+        max_tokens = typer.prompt(
+            "Enter the model's max output tokens", default=8192, 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_id": "gpt-4o",  # This is a placeholder to use the Azure PromptEngine
+            "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.llm.models_info import MODEL_CONFIG_DIR, 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,134 @@
+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", "-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,
+    use_janus_inputs: Annotated[
+        bool,
+        typer.Option(
+            "-j",
+            "--use-janus-inputs",
+            help="Present if converter should use janus inputs",
+        ),
+    ] = False,
+):
+    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,
+        use_janus_inputs=use_janus_inputs,
+    )
+    partitioner = Partitioner(**kwargs)
+    partitioner.translate(input_dir, output_dir, failure_dir, overwrite)

janus/cli/pipeline.py

@@ -0,0 +1,123 @@
+import json
+from pathlib import Path
+from typing import Optional
+
+import click
+import typer
+from typing_extensions import Annotated
+
+from janus.cli.constants import CONVERTERS
+from janus.converter.chain import ConverterChain
+from janus.utils.enums import LANGUAGES
+
+
+def instiantiate(x):
+    if isinstance(x, dict):
+        if "type" in x:
+            if "args" not in x:
+                x["args"] = []
+            x["args"] = [instiantiate(a) for a in x["args"]]
+            if "kwargs" not in x:
+                x["kwargs"] = {}
+            x["kwargs"] = {k: instiantiate(x["kwargs"][k]) for k in x["kwargs"]}
+            if x["type"] not in CONVERTERS:
+                raise ValueError(f"Error: {x['type']} is not a Converter")
+            return CONVERTERS[x["type"]](*x["args"], **x["kwargs"])
+        else:
+            return {k: instiantiate(x[k]) for k in x}
+    elif isinstance(x, list):
+        return [instiantiate(a) for a in x]
+    else:
+        return x
+
+
+def instiantiate_pipeline(
+    pipeline: list[dict],
+    language: str = "text",
+    model: str = "gpt-4o",
+    use_janus_inputs: None | bool = None,
+):
+    if "kwargs" not in pipeline[0]:
+        pipeline[0]["kwargs"] = {}
+    pipeline[0]["kwargs"].update(source_language=language, model=model)
+    if use_janus_inputs is not None:
+        pipeline[0]["kwargs"].update(janus_inputs=use_janus_inputs)
+    print(pipeline[0])
+    converters = [instiantiate(pipeline[0])]
+    for p in pipeline[1:]:
+        p["kwargs"].update(source_language=converters[-1].target_language, model=model)
+        converters.append(instiantiate(p))
+    return ConverterChain(*converters)
+
+
+def pipeline(
+    pipeline_file: Annotated[
+        Path, typer.Option("-p", "--pipeline", help="Name of pipeline file to use")
+    ],
+    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.",
+        ),
+    ],
+    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 translated code in."
+        ),
+    ],
+    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 documentation",
+        ),
+    ] = None,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite/--preserve",
+            help="Whether to overwrite existing files in the output directory",
+        ),
+    ] = False,
+    use_janus_inputs: Annotated[
+        Optional[bool],
+        typer.Option(
+            "-j",
+            "--use-janus-inputs",
+            help="Present if converter chain should use janus input files",
+        ),
+    ] = None,
+):
+    with open(pipeline_file, "r") as f:
+        json_obj = json.load(f)
+    pipeline = instiantiate_pipeline(
+        json_obj, language=language, model=llm_name, use_janus_inputs=use_janus_inputs
+    )
+    pipeline.translate(
+        input_directory=input_dir,
+        output_directory=output_dir,
+        failure_directory=failure_dir,
+        overwrite=overwrite,
+    )