janus-llm 1.0.0__py3-none-any.whl → 2.0.1__py3-none-any.whl

janus/llm/model_callbacks.py

@@ -0,0 +1,177 @@
+import threading
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Generator
+
+from langchain_core.callbacks import BaseCallbackHandler
+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
+
+log = create_logger(__name__)
+
+
+# Updated 2024-06-21
+COST_PER_1K_TOKENS: dict[str, dict[str, float]] = {
+    "gpt-3.5-turbo-0125": {"input": 0.0005, "output": 0.0015},
+    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-0125-preview": {"input": 0.01, "output": 0.03},
+    "gpt-4-0613": {"input": 0.03, "output": 0.06},
+    "gpt-4o-2024-05-13": {"input": 0.005, "output": 0.015},
+    "anthropic.claude-v2": {"input": 0.008, "output": 0.024},
+    "anthropic.claude-instant-v1": {"input": 0.0008, "output": 0.0024},
+    "anthropic.claude-3-haiku-20240307-v1:0": {"input": 0.00025, "output": 0.00125},
+    "anthropic.claude-3-sonnet-20240229-v1:0": {"input": 0.003, "output": 0.015},
+    "meta.llama2-13b-chat-v1": {"input": 0.00075, "output": 0.001},
+    "meta.llama2-70b-chat-v1": {"input": 0.00195, "output": 0.00256},
+    "meta.llama2-13b-v1": {"input": 0.0, "output": 0.0},
+    "meta.llama2-70b-v1": {"input": 0.00265, "output": 0.0035},
+    "meta.llama3-8b-instruct-v1:0": {"input": 0.0003, "output": 0.0006},
+    "meta.llama3-70b-instruct-v1:0": {"input": 0.00265, "output": 0.0035},
+    "amazon.titan-text-lite-v1": {"input": 0.00015, "output": 0.0002},
+    "amazon.titan-text-express-v1": {"input": 0.0002, "output": 0.0006},
+    "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},
+}
+
+
+def _get_token_cost(
+    prompt_tokens: int, completion_tokens: int, model_id: str | None
+) -> float:
+    """Get the cost of tokens according to model ID"""
+    if model_id not in COST_PER_1K_TOKENS:
+        raise ValueError(
+            f"Unknown model: {model_id}. Please provide a valid model name."
+            f" Known models are: {', '.join(COST_PER_1K_TOKENS.keys())}"
+        )
+    model_cost = COST_PER_1K_TOKENS[model_id]
+    input_cost = (prompt_tokens / 1000.0) * model_cost["input"]
+    output_cost = (completion_tokens / 1000.0) * model_cost["output"]
+    return input_cost + output_cost
+
+
+class TokenUsageCallbackHandler(BaseCallbackHandler):
+    """Callback Handler that tracks metadata on model cost, retries, etc.
+    Based on https://github.com/langchain-ai/langchain/blob/master/libs
+                /community/langchain_community/callbacks/openai_info.py
+    """
+
+    total_tokens: int = 0
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    successful_requests: int = 0
+    total_cost: float = 0.0
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._lock = threading.Lock()
+
+    def __repr__(self) -> str:
+        return (
+            f"Tokens Used: {self.total_tokens}\n"
+            f"\tPrompt Tokens: {self.prompt_tokens}\n"
+            f"\tCompletion Tokens: {self.completion_tokens}\n"
+            f"Successful Requests: {self.successful_requests}\n"
+            f"Total Cost (USD): ${self.total_cost}"
+        )
+
+    @property
+    def always_verbose(self) -> bool:
+        """Whether to call verbose callbacks even if verbose is False."""
+        return True
+
+    def on_chat_model_start(self, *args, **kwargs):
+        pass
+
+    def on_llm_start(
+        self, serialized: dict[str, Any], prompts: list[str], **kwargs: Any
+    ) -> None:
+        """Print out the prompts."""
+        pass
+
+    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
+        """Print out the token."""
+        pass
+
+    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
+        """Collect token usage."""
+        # Check for usage_metadata (langchain-core >= 0.2.2)
+        try:
+            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", "")
+
+        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", "")
+
+        total_cost = _get_token_cost(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            model_id=model_name,
+        )
+
+        # update shared state behind lock
+        with self._lock:
+            self.total_cost += total_cost
+            self.total_tokens += token_usage.get("total_tokens", 0)
+            self.prompt_tokens += prompt_tokens
+            self.completion_tokens += completion_tokens
+            self.successful_requests += 1
+
+    def __copy__(self) -> "TokenUsageCallbackHandler":
+        """Return a copy of the callback handler."""
+        return self
+
+    def __deepcopy__(self, memo: Any) -> "TokenUsageCallbackHandler":
+        """Return a deep copy of the callback handler."""
+        return self
+
+
+token_usage_callback_var: ContextVar[TokenUsageCallbackHandler | None] = ContextVar(
+    "token_usage_callback_var", default=None
+)
+register_configure_hook(token_usage_callback_var, True)
+
+
+@contextmanager
+def get_model_callback() -> Generator[TokenUsageCallbackHandler, None, None]:
+    cb = TokenUsageCallbackHandler()
+    token_usage_callback_var.set(cb)
+    yield cb
+    token_usage_callback_var.set(None)

janus/llm/models_info.py

@@ -1,115 +1,208 @@
 import json
 import os
 from pathlib import Path
-from typing import Any, Dict, Tuple
+from typing import Any, Callable
 
 from dotenv import load_dotenv
-from langchain.chat_models import ChatOpenAI
-from langchain.llms import HuggingFaceTextGenInference
-from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
+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 (
+    ChatGptPromptEngine,
+    ClaudePromptEngine,
+    CoherePromptEngine,
+    Llama2PromptEngine,
+    Llama3PromptEngine,
+    PromptEngine,
+    TitanPromptEngine,
+)
+
+from ..utils.logger import create_logger
+
+log = create_logger(__name__)
+
+try:
+    from langchain_community.chat_models import BedrockChat
+    from langchain_community.llms.bedrock import Bedrock
+except ImportError:
+    log.warning(
+        "Could not import LangChain's Bedrock Client. If you would like to use Bedrock "
+        "models, please install LangChain's Bedrock Client by running 'pip install "
+        "janus-llm[bedrock]' or poetry install -E bedrock."
+    )
+
+try:
+    from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
+except ImportError:
+    log.warning(
+        "Could not import LangChain's HuggingFace Pipeline Client. If you would like to "
+        "use HuggingFace models, please install LangChain's HuggingFace Pipeline Client "
+        "by running 'pip install janus-llm[hf-local]' or poetry install -E hf-local."
+    )
+
 
 load_dotenv()
 
-MODEL_TYPE_CONSTRUCTORS = {
+openai_model_reroutes = {
+    "gpt-4o": "gpt-4o-2024-05-13",
+    "gpt-4": "gpt-4-0613",
+    "gpt-4-turbo": "gpt-4-turbo-2024-04-09",
+    "gpt-4-turbo-preview": "gpt-4-0125-preview",
+    "gpt-3.5-turbo": "gpt-3.5-turbo-0125",
+    "gpt-3.5-turbo-16k": "gpt-3.5-turbo-0125",
+}
+
+openai_models = [
+    "gpt-4-0613",
+    "gpt-4-1106-preview",
+    "gpt-4-0125-preview",
+    "gpt-4o-2024-05-13",
+    "gpt-3.5-turbo-0125",
+]
+claude_models = [
+    "bedrock-claude-v2",
+    "bedrock-claude-instant-v1",
+    "bedrock-claude-haiku",
+    "bedrock-claude-sonnet",
+]
+llama2_models = [
+    "bedrock-llama2-70b",
+    "bedrock-llama2-70b-chat",
+    "bedrock-llama2-13b",
+    "bedrock-llama2-13b-chat",
+]
+llama3_models = [
+    "bedrock-llama3-8b-instruct",
+    "bedrock-llama3-70b-instruct",
+]
+titan_models = [
+    "bedrock-titan-text-lite",
+    "bedrock-titan-text-express",
+    "bedrock-jurassic-2-mid",
+    "bedrock-jurassic-2-ultra",
+]
+cohere_models = [
+    "bedrock-command-r-plus",
+]
+bedrock_models = [
+    *claude_models,
+    *llama2_models,
+    *llama3_models,
+    *titan_models,
+    *cohere_models,
+]
+all_models = [*openai_models, *bedrock_models]
+
+MODEL_TYPE_CONSTRUCTORS: dict[str, Callable[[Any], BaseLanguageModel]] = {
     "OpenAI": ChatOpenAI,
     "HuggingFace": HuggingFaceTextGenInference,
-    "HuggingFaceLocal": HuggingFacePipeline.from_model_id,
 }
 
+try:
+    MODEL_TYPE_CONSTRUCTORS.update(
+        {
+            "HuggingFaceLocal": HuggingFacePipeline.from_model_id,
+            "Bedrock": Bedrock,
+            "BedrockChat": BedrockChat,
+        }
+    )
+except NameError:
+    pass
+
 
-MODEL_TYPES: Dict[str, Any] = {
-    "gpt-4": "OpenAI",
-    "gpt-4-32k": "OpenAI",
-    "gpt-4-1106-preview": "OpenAI",
-    "gpt-3.5-turbo": "OpenAI",
-    "gpt-3.5-turbo-16k": "OpenAI",
-    "mitre-llama": "HuggingFace",
-    "mitre-falcon": "HuggingFace",
-    "mitre-wizard-coder": "HuggingFace",
+MODEL_PROMPT_ENGINES: dict[str, Callable[..., PromptEngine]] = {
+    **{m: ChatGptPromptEngine for m in openai_models},
+    **{m: ClaudePromptEngine for m in claude_models},
+    **{m: Llama2PromptEngine for m in llama2_models},
+    **{m: Llama3PromptEngine for m in llama3_models},
+    **{m: TitanPromptEngine for m in titan_models},
+    **{m: CoherePromptEngine for m in cohere_models},
 }
 
-_open_ai_defaults: Dict[str, Any] = {
+_open_ai_defaults: dict[str, str] = {
     "openai_api_key": os.getenv("OPENAI_API_KEY"),
     "openai_organization": os.getenv("OPENAI_ORG_ID"),
 }
 
-MODEL_DEFAULT_ARGUMENTS: Dict[str, Dict[str, Any]] = {
-    "gpt-4": dict(model_name="gpt-4"),
-    "gpt-4-32k": dict(model_name="gpt-4-32k"),
-    "gpt-4-1106-preview": dict(model_name="gpt-4-1106-preview"),
-    "gpt-3.5-turbo": dict(model_name="gpt-3.5-turbo"),
-    "gpt-3.5-turbo-16k": dict(model_name="gpt-3.5-turbo-16k"),
-    "mitre-llama": dict(
-        inference_server_url="https://llama2-70b.aip.mitre.org",
-        max_new_tokens=4096,
-        top_k=10,
-        top_p=0.95,
-        typical_p=0.95,
-        temperature=0.01,
-        repetition_penalty=1.03,
-        timeout=240,
-    ),
-    "mitre-falcon": dict(
-        inference_server_url="https://falcon-40b.aip.mitre.org",
-        max_new_tokens=4096,
-        top_k=10,
-        top_p=0.95,
-        typical_p=0.95,
-        temperature=0.01,
-        repetition_penalty=1.03,
-        timeout=240,
-    ),
-    "mitre-wizard-coder": dict(
-        inference_server_url="https://wizard-coder-34b.aip.mitre.org",
-        max_new_tokens=1024,
-        top_k=10,
-        top_p=0.95,
-        typical_p=0.95,
-        temperature=0.01,
-        repetition_penalty=1.03,
-        timeout=240,
-    ),
+model_identifiers = {
+    **{m: m for m in openai_models},
+    "bedrock-claude-v2": "anthropic.claude-v2",
+    "bedrock-claude-instant-v1": "anthropic.claude-instant-v1",
+    "bedrock-claude-haiku": "anthropic.claude-3-haiku-20240307-v1:0",
+    "bedrock-claude-sonnet": "anthropic.claude-3-sonnet-20240229-v1:0",
+    "bedrock-llama2-70b": "meta.llama2-70b-v1",
+    "bedrock-llama2-70b-chat": "meta.llama2-70b-chat-v1",
+    "bedrock-llama2-13b": "meta.llama2-13b-chat-v1",
+    "bedrock-llama2-13b-chat": "meta.llama2-13b-v1",
+    "bedrock-llama3-8b-instruct": "meta.llama3-8b-instruct-v1:0",
+    "bedrock-llama3-70b-instruct": "meta.llama3-70b-instruct-v1:0",
+    "bedrock-titan-text-lite": "amazon.titan-text-lite-v1",
+    "bedrock-titan-text-express": "amazon.titan-text-express-v1",
+    "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",
+}
+
+MODEL_DEFAULT_ARGUMENTS: dict[str, dict[str, str]] = {
+    k: (dict(model_name=k) if k in openai_models else dict(model_id=v))
+    for k, v in model_identifiers.items()
 }
 
 DEFAULT_MODELS = list(MODEL_DEFAULT_ARGUMENTS.keys())
 
 MODEL_CONFIG_DIR = Path.home().expanduser() / ".janus" / "llm"
 
-TOKEN_LIMITS: Dict[str, int] = {
-    "gpt-4": 8192,
+MODEL_TYPES: dict[str, PromptEngine] = {
+    **{m: "OpenAI" for m in openai_models},
+    **{m: "BedrockChat" for m in bedrock_models},
+}
+
+TOKEN_LIMITS: dict[str, int] = {
     "gpt-4-32k": 32_768,
+    "gpt-4-0613": 8192,
     "gpt-4-1106-preview": 128_000,
-    "gpt-3.5-turbo": 4096,
-    "gpt-3.5-turbo-16k": 16_384,
-    "mitre-falcon": 32_000,
+    "gpt-4-0125-preview": 128_000,
+    "gpt-4o-2024-05-13": 128_000,
+    "gpt-3.5-turbo-0125": 16_384,
     "text-embedding-ada-002": 8191,
     "gpt4all": 16_384,
-}
-
-COST_PER_MODEL: Dict[str, Dict[str, float]] = {
-    "gpt-4": {"input": 0.03, "output": 0.06},
-    "gpt-4-32k": {"input": 0.6, "output": 0.12},
-    "gpt-4-1106-preview": {"input": 0.01, "output": 0.03},
-    "gpt-3.5-turbo": {"input": 0.0015, "output": 0.002},
-    "gpt-3.5-turbo-16k": {"input": 0.003, "output": 0.004},
-    "mitre-llama": {"input": 0.0, "output": 0.0},
-    "mitre-falcon": {"input": 0.0, "output": 0.0},
-    "mitre-wizard-coder": {"input": 0.0, "output": 0.0},
+    "anthropic.claude-v2": 100_000,
+    "anthropic.claude-instant-v1": 100_000,
+    "anthropic.claude-3-haiku-20240307-v1:0": 248_000,
+    "anthropic.claude-3-sonnet-20240229-v1:0": 248_000,
+    "meta.llama2-70b-v1": 4096,
+    "meta.llama2-70b-chat-v1": 4096,
+    "meta.llama2-13b-chat-v1": 4096,
+    "meta.llama2-13b-v1": 4096,
+    "meta.llama3-8b-instruct-v1:0": 8000,
+    "meta.llama3-70b-instruct-v1:0": 8000,
+    "amazon.titan-text-lite-v1": 4096,
+    "amazon.titan-text-express-v1": 8192,
+    "ai21.j2-mid-v1": 8192,
+    "ai21.j2-ultra-v1": 8192,
+    "cohere.command-r-plus-v1:0": 128_000,
 }
 
 
-def load_model(model_name: str) -> Tuple[Any, int, Dict[str, float]]:
+def load_model(model_name: str) -> tuple[BaseLanguageModel, int, dict[str, float]]:
     if not MODEL_CONFIG_DIR.exists():
         MODEL_CONFIG_DIR.mkdir(parents=True)
     model_config_file = MODEL_CONFIG_DIR / f"{model_name}.json"
     if not model_config_file.exists():
         if model_name not in DEFAULT_MODELS:
-            raise ValueError(f"Error: could not find model {model_name}")
+            if model_name in openai_model_reroutes:
+                model_name = openai_model_reroutes[model_name]
+            else:
+                raise ValueError(f"Error: could not find model {model_name}")
         model_config = {
             "model_type": MODEL_TYPES[model_name],
             "model_args": MODEL_DEFAULT_ARGUMENTS[model_name],
-            "token_limit": TOKEN_LIMITS.get(model_name, 4096),
-            "model_cost": COST_PER_MODEL.get(model_name, {"input": 0, "output": 0}),
+            "token_limit": TOKEN_LIMITS.get(model_identifiers[model_name], 4096),
+            "model_cost": COST_PER_1K_TOKENS.get(
+                model_identifiers[model_name], {"input": 0, "output": 0}
+            ),
         }
         with open(model_config_file, "w") as f:
             json.dump(model_config, f)

janus/metrics/__init__.py

@@ -0,0 +1,8 @@
+import glob
+import os.path
+
+modules = glob.glob(os.path.join(os.path.dirname(__file__), "*.py"))
+__all__ = []
+for m in modules:
+    if os.path.isfile(m) and not os.path.samefile(m, __file__):
+        __all__.append(os.path.basename(m)[:-3])

janus/metrics/_tests/reference.py

@@ -0,0 +1,2 @@
+# Hello
+pass

janus/metrics/_tests/target.py

@@ -0,0 +1,2 @@
+# Hello
+pass

janus/metrics/_tests/test_bleu.py

@@ -0,0 +1,56 @@
+import unittest
+
+from sacrebleu import sentence_bleu
+
+from ..bleu import bleu
+
+
+class TestBLEU(unittest.TestCase):
+    def setUp(self):
+        self.target_text = "This is a source text."
+        self.reference_text = "This is a destination text."
+
+    def test_bleu(self):
+        """Test the BLEU score calculation."""
+        function_score = (
+            sentence_bleu(
+                self.target_text,
+                [self.reference_text],
+            ).score
+            / 100.0
+        )
+        expected_score = bleu(
+            self.target_text,
+            self.reference_text,
+        )
+        self.assertEqual(function_score, expected_score)
+
+    def test_bleu_with_s_flag(self):
+        """Test the BLEU score calculation with the -S flag."""
+        function_score = (
+            sentence_bleu(
+                self.target_text,
+                [self.reference_text],
+            ).score
+            / 100.0
+        )
+        score_with_s_flag = bleu(
+            self.target_text,
+            self.reference_text,
+            use_strings=True,  # Mimics -S
+        )
+        self.assertEqual(function_score, score_with_s_flag)
+
+    def test_bleu_invalid_target_type(self):
+        """Test the BLEU score calculation with invalid source text type."""
+        with self.assertRaises(TypeError):
+            sentence_bleu(123, [self.reference_text])
+
+    def test_bleu_invalid_reference_type(self):
+        """Test the BLEU score calculation with invalid destination text type."""
+        with self.assertRaises(TypeError):
+            sentence_bleu(self.target_text, 123)
+
+
+if __name__ == "__main__":
+    unittest.main()

janus/metrics/_tests/test_chrf.py

@@ -0,0 +1,67 @@
+import unittest
+
+from sacrebleu import sentence_chrf
+
+from ..chrf import chrf
+
+
+class TestChrF(unittest.TestCase):
+    def setUp(self):
+        self.target_text = "This is a source text."
+        self.reference_text = "This is a destination text."
+        self.char_order = 6
+        self.word_order = 2
+        self.beta = 2.0
+
+    def test_chrf_custom_params(self):
+        """Test the chrf function with custom parameters."""
+        function_score = chrf(
+            self.target_text,
+            self.reference_text,
+            self.char_order,
+            self.word_order,
+            self.beta,
+        )
+        score = sentence_chrf(
+            hypothesis=self.target_text,
+            references=[self.reference_text],
+            char_order=self.char_order,
+            word_order=self.word_order,
+            beta=self.beta,
+        )
+        expected_score = float(score.score) / 100.0
+        self.assertEqual(function_score, expected_score)
+
+    def test_chrf_with_s_flag(self):
+        """Test the CHRF score calculation with the -S flag."""
+        function_score = sentence_chrf(
+            hypothesis=self.target_text,
+            references=[self.reference_text],
+            char_order=self.char_order,
+            word_order=self.word_order,
+            beta=self.beta,
+        )
+        function_score = float(function_score.score) / 100.0
+        score_with_s_flag = chrf(
+            self.target_text,
+            self.reference_text,
+            self.char_order,
+            self.word_order,
+            self.beta,
+            use_strings=True,  # Mimics -S
+        )
+        self.assertEqual(function_score, score_with_s_flag)
+
+    def test_chrf_invalid_target_type(self):
+        """Test the chrf function with invalid source text type."""
+        with self.assertRaises(TypeError):
+            chrf(123, self.reference_text, self.char_order, self.word_order, self.beta)
+
+    def test_chrf_invalid_reference_type(self):
+        """Test the chrf function with invalid destination text type."""
+        with self.assertRaises(TypeError):
+            chrf(self.target_text, 123, self.char_order, self.word_order, self.beta)
+
+
+if __name__ == "__main__":
+    unittest.main()

janus/metrics/_tests/test_file_pairing.py

@@ -0,0 +1,59 @@
+# FILEPATH: /Users/mdoyle/projects/janus/janus/metrics/tests/test_file_pairing.py
+
+import unittest
+from pathlib import Path
+
+from ..file_pairing import (
+    FILE_PAIRING_METHODS,
+    pair_by_file,
+    pair_by_line,
+    pair_by_line_comment,
+    register_pairing_method,
+)
+
+
+class TestFilePairing(unittest.TestCase):
+    def setUp(self):
+        self.src = "Hello\nWorld"
+        self.cmp = "Hello\nPython"
+        self.state = {
+            "token_limit": 100,
+            "llm": None,
+            "lang": "python",
+            "target_file": self.src,
+            "cmp_file": self.cmp,
+        }
+
+    def test_register_pairing_method(self):
+        @register_pairing_method(name="test")
+        def test_method(src, cmp, state):
+            return [(src, cmp)]
+
+        self.assertIn("test", FILE_PAIRING_METHODS)
+
+    def test_pair_by_file(self):
+        expected = [(self.src, self.cmp)]
+        result = pair_by_file(self.src, self.cmp)
+        self.assertEqual(result, expected)
+
+    def test_pair_by_line(self):
+        expected = [("Hello", "Hello"), ("World", "Python")]
+        result = pair_by_line(self.src, self.cmp)
+        self.assertEqual(result, expected)
+
+    def test_pair_by_line_comment(self):
+        # This test assumes that the source and comparison files have comments on the
+        # same lines
+        # You may need to adjust this test based on your specific use case
+        self.target = Path(__file__).parent / "target.py"
+        self.reference = Path(__file__).parent / "reference.py"
+        kwargs = {
+            "token_limit": 100,
+            "llm": None,
+            "lang": "python",
+            "target_file": self.target,
+            "reference_file": self.reference,
+        }
+        expected = [("# Hello\n", "# Hello\n")]
+        result = pair_by_line_comment(self.src, self.cmp, **kwargs)
+        self.assertEqual(result, expected)