janus-llm 2.0.2__py3-none-any.whl → 3.0.0__py3-none-any.whl

janus/language/naive/simple_ast.py

@@ -1,18 +1,29 @@
-from janus.language.naive.registry import register_splitter
-from janus.language.treesitter import TreeSitterSplitter
-from janus.utils.enums import LANGUAGES
+from ...utils.enums import LANGUAGES
+from ..alc.alc import AlcSplitter
+from ..mumps.mumps import MumpsSplitter
+from ..treesitter import TreeSitterSplitter
+from .registry import register_splitter
 
 
 @register_splitter("ast-flex")
-class FlexibleTreeSitterSplitter(TreeSitterSplitter):
-    pass
+def get_flexible_ast(language: str, **kwargs):
+    if language == "ibmhlasm":
+        return AlcSplitter(**kwargs)
+    elif language == "mumps":
+        return MumpsSplitter(**kwargs)
+    else:
+        return TreeSitterSplitter(language=language, **kwargs)
 
 
 @register_splitter("ast-strict")
-class StrictTreeSitterSplitter(TreeSitterSplitter):
-    def __init__(self, language: str, **kwargs):
-        kwargs.update(
-            protected_node_types=(LANGUAGES[language]["functional_node_type"],),
-            prune_unprotected=True,
-        )
-        super().__init__(language=language, **kwargs)
+def get_strict_ast(language: str, **kwargs):
+    kwargs.update(
+        protected_node_types=LANGUAGES[language]["functional_node_types"],
+        prune_unprotected=True,
+    )
+    if language == "ibmhlasm":
+        return AlcSplitter(**kwargs)
+    elif language == "mumps":
+        return MumpsSplitter(**kwargs)
+    else:
+        return TreeSitterSplitter(language=language, **kwargs)

janus/language/naive/tag_splitter.py

@@ -1,7 +1,7 @@
-from janus.language.block import CodeBlock
-from janus.language.naive.registry import register_splitter
-from janus.language.node import NodeType
-from janus.language.splitter import Splitter
+from ..block import CodeBlock
+from ..node import NodeType
+from ..splitter import Splitter
+from .registry import register_splitter
 
 
 @register_splitter("tag")

janus/language/splitter.py

@@ -47,8 +47,8 @@ class Splitter(FileManager):
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
         skip_merge: bool = False,
-        protected_node_types: tuple[str] = (),
-        prune_node_types: tuple[str] = (),
+        protected_node_types: tuple[str, ...] = (),
+        prune_node_types: tuple[str, ...] = (),
         prune_unprotected: bool = False,
     ):
         """
@@ -340,7 +340,10 @@ class Splitter(FileManager):
         # Double check length (in theory this should never be an issue)
         tokens = self._count_tokens(text)
         if tokens > self.max_tokens:
-            log.error(f"Merged node ({name}) too long for context!")
+            log.error(
+                f"Merged node ({name}) too long for context!"
+                f" ({tokens} > {self.max_tokens})"
+            )
 
         return CodeBlock(
             text=text,
@@ -420,7 +423,10 @@ class Splitter(FileManager):
             name = f"{node.name}-L#{node_line}"
             tokens = self._count_tokens(line)
             if tokens > self.max_tokens:
-                raise TokenLimitError(r"Irreducible node too large for context!")
+                raise TokenLimitError(
+                    "Irreducible node too large for context!"
+                    f" ({tokens} > {self.max_tokens})"
+                )
 
             node.children.append(
                 CodeBlock(

janus/language/treesitter/treesitter.py

@@ -1,6 +1,7 @@
 import os
 import platform
 from collections import defaultdict
+from ctypes import c_void_p, cdll
 from pathlib import Path
 from typing import Optional
 
@@ -26,8 +27,8 @@ class TreeSitterSplitter(Splitter):
         language: str,
         model: None | BaseLanguageModel = None,
         max_tokens: int = 4096,
-        protected_node_types: tuple[str] = (),
-        prune_node_types: tuple[str] = (),
+        protected_node_types: tuple[str, ...] = (),
+        prune_node_types: tuple[str, ...] = (),
         prune_unprotected: bool = False,
     ) -> None:
         """Initialize a TreeSitterSplitter instance.
@@ -48,10 +49,10 @@ class TreeSitterSplitter(Splitter):
         self._load_parser()
 
     def _get_ast(self, code: str) -> CodeBlock:
-        code = bytes(code, "utf-8")
-        tree = self.parser.parse(code)
+        code_bytes = bytes(code, "utf-8")
+        tree = self.parser.parse(code_bytes)
         root = tree.walk().node
-        root = self._node_to_block(root, code)
+        root = self._node_to_block(root, code_bytes)
         return root
 
     # Recursively print tree to view parsed output (dev helper function)
@@ -98,7 +99,7 @@ class TreeSitterSplitter(Splitter):
 
         text = node.text.decode()
         children = [self._node_to_block(child, original_text) for child in node.children]
-        node = CodeBlock(
+        return CodeBlock(
             id=node.id,
             name=str(node.id),
             text=text,
@@ -112,7 +113,6 @@ class TreeSitterSplitter(Splitter):
             language=self.language,
             tokens=self._count_tokens(text),
         )
-        return node
 
     def _load_parser(self) -> None:
         """Load the parser for the given language.
@@ -139,7 +139,25 @@ class TreeSitterSplitter(Splitter):
 
         # Load the parser using the generated .so file
         self.parser: tree_sitter.Parser = tree_sitter.Parser()
-        self.parser.set_language(tree_sitter.Language(so_file, self.language))
+        pointer = self._so_to_pointer(so_file)
+        self.parser.set_language(tree_sitter.Language(pointer, self.language))
+
+    def _so_to_pointer(self, so_file: str) -> int:
+        """Convert the .so file to a pointer.
+
+        Taken from `treesitter.Language.__init__` to get past deprecated warning.
+
+        Arguments:
+            so_file: The path to the so file for the language.
+
+        Returns:
+            The pointer to the language.
+        """
+        lib = cdll.LoadLibrary(os.fspath(so_file))
+        language_function = getattr(lib, f"tree_sitter_{self.language}")
+        language_function.restype = c_void_p
+        pointer = language_function()
+        return pointer
 
     def _create_parser(self, so_file: Path | str) -> None:
         """Create the parser for the given language.

janus/llm/model_callbacks.py

@@ -8,7 +8,7 @@ from langchain_core.messages import AIMessage
 from langchain_core.outputs import ChatGeneration, LLMResult
 from langchain_core.tracers.context import register_configure_hook
 
-from janus.utils.logger import create_logger
+from ..utils.logger import create_logger
 
 log = create_logger(__name__)
 
@@ -35,6 +35,9 @@ COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
     "ai21.j2-mid-v1": {"input": 0.0125, "output": 0.0125},
     "ai21.j2-ultra-v1": {"input": 0.0188, "output": 0.0188},
     "cohere.command-r-plus-v1:0": {"input": 0.003, "output": 0.015},
+    "mistral.mistral-7b-instruct-v0:2": {"input": 0.00015, "output": 0.0002},
+    "mistral.mixtral-8x7b-instruct-v0:1": {"input": 0.00045, "output": 0.0007},
+    "mistral.mistral-large-2402-v1:0": {"input": 0.004, "output": 0.012},
 }
 
 
@@ -103,53 +106,47 @@ class TokenUsageCallbackHandler(BaseCallbackHandler):
             generation = response.generations[0][0]
         except IndexError:
             generation = None
-        if isinstance(generation, ChatGeneration):
-            try:
-                message = generation.message
-                if isinstance(message, AIMessage):
-                    usage_metadata = message.usage_metadata
-                else:
-                    usage_metadata = None
-            except AttributeError:
-                usage_metadata = None
-        else:
-            usage_metadata = None
-        if usage_metadata:
-            token_usage = {"total_tokens": usage_metadata["total_tokens"]}
-            completion_tokens = usage_metadata["output_tokens"]
-            prompt_tokens = usage_metadata["input_tokens"]
-            if response.llm_output is None:
-                # model name (and therefore cost) is unavailable in
-                # streaming responses
-                model_name = ""
-            else:
-                model_name = response.llm_output.get("model_name", "")
 
+        model_id = ""
+        usage_metadata = None
+        if hasattr(response, "llm_output") and response.llm_output is not None:
+            model_id = response.llm_output.get("model_id", model_id)
+            model_id = response.llm_output.get("model_name", model_id)
+            usage_metadata = response.llm_output.get("usage", usage_metadata)
+            usage_metadata = response.llm_output.get("token_usage", usage_metadata)
+        elif isinstance(generation, ChatGeneration):
+            if hasattr(generation, "response_metadata"):
+                model_id = generation.response_metadata.get("model_id", model_id)
+                model_id = generation.response_metadata.get("model_name", model_id)
+                usage_metadata = generation.response_metadata.get("usage", usage_metadata)
+            elif hasattr(generation, "message"):
+                if isinstance(generation.message, AIMessage):
+                    usage_metadata = generation.message.usage_metadata
+
+        completion_tokens = 0
+        prompt_tokens = 0
+        total_tokens = 0
+        if usage_metadata:
+            prompt_tokens = usage_metadata.get("prompt_tokens", prompt_tokens)
+            prompt_tokens = usage_metadata.get("input_tokens", prompt_tokens)
+            completion_tokens = usage_metadata.get("completion_tokens", completion_tokens)
+            completion_tokens = usage_metadata.get("output_tokens", completion_tokens)
+            total_tokens = usage_metadata.get("total_tokens", total_tokens)
         else:
-            if response.llm_output is None:
-                return None
-
-            if "token_usage" not in response.llm_output:
-                with self._lock:
-                    self.successful_requests += 1
-                return None
-
-            # compute tokens and cost for this request
-            token_usage = response.llm_output["token_usage"]
-            completion_tokens = token_usage.get("completion_tokens", 0)
-            prompt_tokens = token_usage.get("prompt_tokens", 0)
-            model_name = response.llm_output.get("model_name", "")
+            with self._lock:
+                self.successful_requests += 1
+            return None
 
         total_cost = _get_token_cost(
             prompt_tokens=prompt_tokens,
             completion_tokens=completion_tokens,
-            model_id=model_name,
+            model_id=model_id,
         )
 
         # update shared state behind lock
         with self._lock:
             self.total_cost += total_cost
-            self.total_tokens += token_usage.get("total_tokens", 0)
+            self.total_tokens += total_tokens
             self.prompt_tokens += prompt_tokens
             self.completion_tokens += completion_tokens
             self.successful_requests += 1

janus/llm/models_info.py

@@ -8,18 +8,18 @@ from langchain_community.llms import HuggingFaceTextGenInference
 from langchain_core.language_models import BaseLanguageModel
 from langchain_openai import ChatOpenAI
 
-from janus.llm.model_callbacks import COST_PER_1K_TOKENS
-from janus.prompts.prompt import (
+from ..prompts.prompt import (
     ChatGptPromptEngine,
     ClaudePromptEngine,
     CoherePromptEngine,
     Llama2PromptEngine,
     Llama3PromptEngine,
+    MistralPromptEngine,
     PromptEngine,
     TitanPromptEngine,
 )
-
 from ..utils.logger import create_logger
+from .model_callbacks import COST_PER_1K_TOKENS
 
 log = create_logger(__name__)
 
@@ -86,12 +86,18 @@ titan_models = [
 cohere_models = [
     "bedrock-command-r-plus",
 ]
+mistral_models = [
+    "bedrock-mistral-7b-instruct",
+    "bedrock-mistral-large",
+    "bedrock-mixtral",
+]
 bedrock_models = [
     *claude_models,
     *llama2_models,
     *llama3_models,
     *titan_models,
     *cohere_models,
+    *mistral_models,
 ]
 all_models = [*openai_models, *bedrock_models]
 
@@ -119,6 +125,7 @@ MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
     **{m: Llama3PromptEngine for m in llama3_models},
     **{m: TitanPromptEngine for m in titan_models},
     **{m: CoherePromptEngine for m in cohere_models},
+    **{m: MistralPromptEngine for m in mistral_models},
 }
 
 _open_ai_defaults: dict[str, str] = {
@@ -143,6 +150,9 @@ model_identifiers = {
     "bedrock-jurassic-2-mid": "ai21.j2-mid-v1",
     "bedrock-jurassic-2-ultra": "ai21.j2-ultra-v1",
     "bedrock-command-r-plus": "cohere.command-r-plus-v1:0",
+    "bedrock-mixtral": "mistral.mixtral-8x7b-instruct-v0:1",
+    "bedrock-mistral-7b-instruct": "mistral.mistral-7b-instruct-v0:2",
+    "bedrock-mistral-large": "mistral.mistral-large-2402-v1:0",
 }
 
 MODEL_DEFAULT_ARGUMENTS: dict[str, dict[str, str]] = {
@@ -183,6 +193,9 @@ TOKEN_LIMITS: dict[str, int] = {
     "ai21.j2-mid-v1": 8192,
     "ai21.j2-ultra-v1": 8192,
     "cohere.command-r-plus-v1:0": 128_000,
+    "mistral.mixtral-8x7b-instruct-v0:1": 32_000,
+    "mistral.mistral-7b-instruct-v0:2": 32_000,
+    "mistral.mistral-large-2402-v1:0": 32_000,
 }
 
 

janus/metrics/_tests/test_llm.py

@@ -3,8 +3,7 @@ from unittest.mock import patch
 
 import pytest
 
-from janus.llm.models_info import load_model
-
+from ...llm.models_info import load_model
 from ..llm_metrics import llm_evaluate_option, llm_evaluate_ref_option
 
 
@@ -40,7 +39,7 @@ class TestLLMMetrics(unittest.TestCase):
         print("'Hello, world!")
         """
 
-    @patch("janus.llm.models_info.load_model")
+    @patch(".llm.models_info.load_model")
     @patch("janus.metrics.llm_metrics.llm_evaluate")
     @pytest.mark.llm_eval
     def test_llm_self_eval_quality(self, mock_llm_evaluate, mock_load_model):

janus/metrics/_tests/test_rouge_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from janus.metrics.rouge_score import rouge
+from ..rouge_score import rouge
 
 
 class TestRouge(unittest.TestCase):

janus/metrics/_tests/test_similarity_score.py

@@ -1,6 +1,6 @@
 import unittest
 
-from janus.metrics.similarity import similarity_score
+from ..similarity import similarity_score
 
 
 class TestSimilarityScore(unittest.TestCase):

janus/metrics/complexity_metrics.py

@@ -1,10 +1,9 @@
 import math
 from typing import List, Optional
 
-from janus.language.block import CodeBlock
-from janus.language.treesitter.treesitter import TreeSitterSplitter
-from janus.utils.enums import LANGUAGES
-
+from ..language.block import CodeBlock
+from ..language.treesitter.treesitter import TreeSitterSplitter
+from ..utils.enums import LANGUAGES
 from .metric import metric
 
 

janus/metrics/metric.py

@@ -7,10 +7,9 @@ import click
 import typer
 from typing_extensions import Annotated
 
-from janus.llm import load_model
-from janus.utils.enums import LANGUAGES
-from janus.utils.logger import create_logger
-
+from ..llm import load_model
+from ..utils.enums import LANGUAGES
+from ..utils.logger import create_logger
 from ..utils.progress import track
 from .cli import evaluate
 from .file_pairing import FILE_PAIRING_METHODS

janus/metrics/reading.py

@@ -1,9 +1,30 @@
+import re
+
 import nltk
 import readability
+from nltk.tokenize import TweetTokenizer
 
 from .metric import metric
 
 
+def word_count(text):
+    """Calculates word count exactly how readability package does
+
+    Arguments:
+        text: The input string.
+
+    Returns:
+        Word Count
+    """
+    tokenizer = TweetTokenizer()
+    word_count = 0
+    tokens = tokenizer.tokenize(text)
+    for t in tokens:
+        if not re.match(r"^[.,\/#!$%'\^&\*;:{}=\-_`~()]+$", t):
+            word_count += 1
+    return word_count
+
+
 def _repeat_text(text):
     """Repeats a string until its length is over 100 words.
 
@@ -20,11 +41,10 @@ def _repeat_text(text):
     if not text.endswith("."):
         text += "."  # Add a period if missing
 
-    # Check if repeated text is long enough, repeat more if needed
     repeated_text = text
-    while len(repeated_text.split()) < 100:
-        repeated_text += " " + text
 
+    while word_count(repeated_text) < 100:
+        repeated_text += " " + text
     return repeated_text
 
 
@@ -52,7 +72,8 @@ def flesch(target: str, **kwargs) -> float:
     Returns:
         The Flesch score.
     """
-
+    if not target.strip():  # Check if the target text is blank
+        return None
     return get_readability(target).flesch().score
 
 
@@ -66,5 +87,6 @@ def gunning_fog(target: str, **kwargs) -> float:
     Returns:
         The Gunning-Fog score.
     """
-
+    if not target.strip():  # Check if the target text is blank
+        return None
     return get_readability(target).gunning_fog().score

janus/prompts/prompt.py

@@ -2,12 +2,12 @@ import json
 from abc import ABC, abstractmethod
 from pathlib import Path
 
-from langchain import PromptTemplate
 from langchain.prompts import ChatPromptTemplate
 from langchain.prompts.chat import (
     HumanMessagePromptTemplate,
     SystemMessagePromptTemplate,
 )
+from langchain_core.prompts import PromptTemplate
 
 from ..utils.enums import LANGUAGES
 from ..utils.logger import create_logger
@@ -34,15 +34,59 @@ HUMAN_PROMPT_TEMPLATE_FILENAME = "human.txt"
 PROMPT_VARIABLES_FILENAME = "variables.json"
 
 
+retry_with_output_prompt_text = """Instructions:
+--------------
+{instructions}
+--------------
+Completion:
+--------------
+{completion}
+--------------
+
+Above, the Completion did not satisfy the constraints given in the Instructions.
+Error:
+--------------
+{error}
+--------------
+
+Please try again. Please only respond with an answer that satisfies the
+constraints laid out in the Instructions:"""
+
+
+retry_with_error_and_output_prompt_text = """Prompt:
+--------------
+{prompt}
+--------------
+Completion:
+--------------
+{completion}
+--------------
+
+Above, the Completion did not satisfy the constraints given in the Prompt.
+Error:
+--------------
+{error}
+--------------
+
+Please try again. Please only respond with an answer that satisfies the
+constraints laid out in the Prompt:"""
+
+
+retry_with_output_prompt = PromptTemplate.from_template(retry_with_output_prompt_text)
+retry_with_error_and_output_prompt = PromptTemplate.from_template(
+    retry_with_error_and_output_prompt_text
+)
+
+
 class PromptEngine(ABC):
     """A class defining prompting schemes for the LLM."""
 
     def __init__(
         self,
         source_language: str,
-        target_language: str,
-        target_version: str,
         prompt_template: str,
+        target_language: str | None = None,
+        target_version: str | None = None,
     ) -> None:
         """Initialize a PromptEngine instance.
 
@@ -63,15 +107,18 @@ class PromptEngine(ABC):
 
         # Define variables to be passed in to the prompt formatter
         source_language = source_language.lower()
-        target_language = target_language.lower()
         self.variables = dict(
             SOURCE_LANGUAGE=source_language,
-            TARGET_LANGUAGE=target_language,
-            TARGET_LANGUAGE_VERSION=str(target_version),
             FILE_SUFFIX=LANGUAGES[source_language]["suffix"],
             SOURCE_CODE_EXAMPLE=LANGUAGES[source_language]["example"],
-            TARGET_CODE_EXAMPLE=LANGUAGES[target_language]["example"],
         )
+        if target_language is not None:
+            target_language = target_language.lower()
+            self.variables.update(
+                TARGET_LANGUAGE=target_language,
+                TARGET_CODE_EXAMPLE=LANGUAGES[target_language]["example"],
+            )
+            self.variables.update(TARGET_LANGUAGE_VERSION=str(target_version))
         variables_path = template_path / PROMPT_VARIABLES_FILENAME
         if variables_path.exists():
             self.variables.update(json.loads(variables_path.read_text()))
@@ -219,3 +266,16 @@ class CoherePromptEngine(PromptEngine):
             f"{human_prompt}"
             f"<|END_OF_TURN_TOKEN|>"
         )
+
+
+class MistralPromptEngine(PromptEngine):
+    def load_prompt_template(self, template_path: Path) -> ChatPromptTemplate:
+        system_prompt_path = template_path / SYSTEM_PROMPT_TEMPLATE_FILENAME
+        system_prompt = system_prompt_path.read_text()
+
+        human_prompt_path = template_path / HUMAN_PROMPT_TEMPLATE_FILENAME
+        human_prompt = human_prompt_path.read_text()
+
+        return PromptTemplate.from_template(
+            f"<s>[INST] {system_prompt} [/INST] </s>[INST] {human_prompt}  [/INST]"
+        )

janus/utils/enums.py

@@ -10,7 +10,7 @@ class EmbeddingType(Enum):
     TARGET = 5  # placeholder embeddings, are these useful for analysis?
 
 
-CUSTOM_SPLITTERS: Set[str] = {"mumps", "binary"}
+CUSTOM_SPLITTERS: Set[str] = {"mumps", "binary", "ibmhlasm"}
 
 LANGUAGES: Dict[str, Dict[str, Any]] = {
     "ada": {
@@ -63,7 +63,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
             '#include <stdio.h>\n\nint main() {\n    printf("Hello, World!\\n");\n'
             "    return 0;\n}\n"
         ),
-        "functional_node_type": "function_definition",
+        "functional_node_types": ["function_definition"],
         "comment_node_type": "comment",
     },
     "capnp": {
@@ -206,7 +206,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "example": (
             "program HelloWorld\n  print *, 'Hello, World!'\nend program HelloWorld\n"
         ),
-        "functional_node_type": "function",
+        "functional_node_types": ["function"],
         "comment_node_type": "comment",
     },
     "gitattributes": {
@@ -300,6 +300,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
                      END   HELLO
         """
         ),
+        "functional_node_types": ["csect", "dsect"],
         "branch_node_types": ["branch_instruction"],
         "operation_node_types": ["operation", "branch_operation"],
         "operand_node_types": ["operands"],
@@ -420,7 +421,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "suffix": "m",
         "url": "https://github.com/janus-llm/tree-sitter-mumps",
         "example": 'WRITE "Hello, World!"',
-        "functional_node_type": "routine_definition",
+        "functional_node_types": ["routine_definition"],
         "comment_node_type": "comment",
         "branch_node_types": ["if_statement"],
         "operation_node_types": [
@@ -512,7 +513,7 @@ LANGUAGES: Dict[str, Dict[str, Any]] = {
         "suffix": "py",
         "url": "https://github.com/tree-sitter/tree-sitter-python",
         "example": "# Hello, World!\nprint('Hello, World!')\n",
-        "functional_node_type": "function_definition",
+        "functional_node_types": ["function_definition"],
         "comment_node_type": "comment",
     },
     "qmljs": {

{janus_llm-2.0.2.dist-info → janus_llm-3.0.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: janus-llm
-Version: 2.0.2
+Version: 3.0.0
 Summary: A transcoding library using LLMs.
 Home-page: https://github.com/janus-llm/janus-llm
 License: Apache 2.0