eval-framework 0.2.0__py3-none-any.whl

eval_framework/llm/huggingface.py

@@ -0,0 +1,332 @@
+import logging
+from collections.abc import Callable, Sequence
+from functools import partial
+from typing import Any
+
+import torch
+from tokenizers import Tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
+
+from eval_framework.llm.base import BaseLLM
+from eval_framework.shared.types import (
+    ConcatCompression,
+    Error,
+    PromptTooLongException,
+    RawCompletion,
+    RawLoglikelihood,
+)
+from eval_framework.tasks.base import Sample
+from eval_framework.tasks.utils import raise_errors
+from eval_framework.utils.constants import RED, RESET
+from template_formatting.formatter import BaseFormatter, ConcatFormatter, HFFormatter, Llama3Formatter, Message
+
+logger = logging.getLogger(__name__)
+
+
+class StopSequenceCriteria(StoppingCriteria):
+    def __init__(self, tokenizer: Tokenizer, stop_sequences: list[str], prompt_token_count: int) -> None:
+        self.tokenizer = tokenizer
+        self.stop_sequences = stop_sequences
+        self.prompt_token_count = prompt_token_count
+        # (relatively weak) upper bound for the number of tokens that
+        # need to be decoded to check for stop sequences
+        self.token_history_length = max(map(len, stop_sequences))
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: Any) -> bool:
+        sequence = input_ids[0].tolist()
+        sequence = sequence[self.prompt_token_count :]
+        if len(sequence) > self.token_history_length:
+            sequence = sequence[-self.token_history_length :]
+        decoded_text = self.tokenizer.decode(sequence, skip_special_tokens=True)
+
+        for stop_sequence in self.stop_sequences:
+            if stop_sequence in decoded_text:
+                return True
+        return False
+
+
+class RepeatedTokenSequenceCriteria(StoppingCriteria):
+    def __init__(self, tokenizer: Tokenizer, completion_start_index: int) -> None:
+        self.tokenizer = tokenizer
+        # Initialize with an empty string to store the last line
+        self.last_line = ""
+        self.completion_start_index = completion_start_index
+        # self.newline_token_id = tokenizer.encode('\n')
+
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs: Any) -> torch.Tensor:
+        # Convert token ids to tokens
+        tokens = self.tokenizer.decode(input_ids[0][self.completion_start_index :])
+
+        # Join tokens to form the current text
+        current_text = "".join(tokens)
+        # Split text into lines
+        lines = current_text.split("\n")
+
+        # Check if the last full line (ignoring the last if it's incomplete) is repeated
+        if len(lines) > 1 and lines[-2] == lines[-1] and not (lines[-1] == "" and lines[-2] == ""):
+            return torch.BoolTensor([True]).to(input_ids.device)  # Stop generation if repeated line is found
+
+        return torch.BoolTensor([False]).to(input_ids.device)
+
+
+class HFLLM(BaseLLM):
+    LLM_NAME: str
+    DEFAULT_FORMATTER: Callable[[], BaseFormatter] | None = None
+    SEQ_LENGTH: int | None = None
+
+    def __init__(self, formatter: BaseFormatter | None = None) -> None:
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.LLM_NAME)
+        self.model = AutoModelForCausalLM.from_pretrained(self.LLM_NAME, device_map="auto")
+        logger.info(f"{RED}[ Model initialized --------------------- {RESET}{self.LLM_NAME} {RED}]{RESET}")
+        self._set_formatter(formatter)
+
+    def _set_formatter(self, formatter: BaseFormatter | None = None) -> None:
+        # if formatter is being set at initialization time, use it
+        if formatter is not None:
+            self._formatter = formatter
+        # if formatter is not being set at initialization time, but DEFAULT_FORMATTER was specified, use it
+        elif self.DEFAULT_FORMATTER is not None:
+            self._formatter = self.DEFAULT_FORMATTER()
+        # if formatter is not being set at initialization time and there is no default formatter,
+        # using HF chat formatter if exists
+        elif self.tokenizer.chat_template is not None:
+            self._formatter = HFFormatter(self.LLM_NAME)
+        # if formatter is not being set at initialization time and there is no default formatter and no chat formatter,
+        # using ConcatFormatter
+        else:
+            raise ValueError("No formatter specified and no default formatter available.")
+
+        logger.info(
+            f"{RED}[ Using default formatter --------------------- {RESET}{self._formatter.__class__.__name__} {RED}]{RESET}"  # noqa: E501
+        )
+
+    def count_tokens(self, text: str, /) -> int:
+        """Count the number of tokens in a string."""
+        return len(self.tokenizer(text, add_special_tokens=False)["input_ids"])
+
+    def generate_from_messages(
+        self,
+        messages: list[Sequence[Message]],
+        stop_sequences: list[str] | None = None,
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> list[RawCompletion]:
+        if temperature is None:
+            effective_temperature = 0.0  # Current default, TODO: refactor to use model's default
+            logger.info(
+                f"Using default temperature value: {effective_temperature} as no custom temperature value was provided"
+            )
+        else:
+            effective_temperature = temperature
+
+        raw_completions = []
+        for single_messages in messages:
+            # format
+            prompt = self._formatter.format(single_messages, output_mode="string")
+            # add_special_tokens would add a second BOS token without explicitly setting it False
+            inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.device)
+
+            prompt_token_count = len(inputs["input_ids"][0])
+            pad_token_id = self.tokenizer.eos_token_id
+
+            # Prepare stopping criteria
+            stopping_criteria = StoppingCriteriaList()
+            if stop_sequences is not None:
+                stopping_criteria.append(StopSequenceCriteria(self.tokenizer, stop_sequences, prompt_token_count))  # type: ignore[attr-defined]
+
+            stopping_criteria.append(  # type: ignore[attr-defined]
+                RepeatedTokenSequenceCriteria(
+                    self.tokenizer,
+                    prompt_token_count,
+                )
+            )
+
+            min_seq_length = min(filter(None, [self.seq_length, self.SEQ_LENGTH]))
+
+            # Calculate the maximum number of tokens to generate
+            max_tokens_to_generate = min_seq_length - prompt_token_count
+            # If max_tokens is specified, use the smaller of the two
+            max_tokens_to_generate = min(filter(None, [max_tokens_to_generate, max_tokens]))
+
+            if max_tokens_to_generate < 1:
+                if raise_errors():
+                    raise PromptTooLongException("Prompt exceeded context size.")
+                raw_completions.append(
+                    RawCompletion(
+                        prompt=prompt,
+                        prompt_sequence_positions=prompt_token_count,
+                        completion="",
+                        completion_sequence_positions=0,
+                        raw_completion_error=Error(
+                            error_class=PromptTooLongException.__name__,
+                            message="Prompt exceeded context size.",
+                            traceback="",
+                        ),
+                    )
+                )
+                continue
+
+            completion, completion_token_count = self._model_generate(
+                redis_key=(prompt, stop_sequences, max_tokens_to_generate, effective_temperature),
+                prompt_token_count=prompt_token_count,
+                inputs=inputs["input_ids"],
+                max_new_tokens=max_tokens_to_generate,
+                stopping_criteria=stopping_criteria,
+                num_return_sequences=1,
+                pad_token_id=pad_token_id,
+                return_dict_in_generate=False,
+                output_scores=False,
+                do_sample=effective_temperature > 0,
+                temperature=effective_temperature if effective_temperature > 0 else None,
+            )
+
+            raw_completions.append(
+                RawCompletion(
+                    prompt=prompt,
+                    prompt_sequence_positions=prompt_token_count,
+                    concat_compression=ConcatCompression.calculate(
+                        single_messages, count_tokens=self.count_tokens, completion=completion
+                    ),
+                    completion=completion,
+                    completion_sequence_positions=completion_token_count,
+                )
+            )
+        return raw_completions
+
+    def _model_generate(self, redis_key: Any, prompt_token_count: int, **kwargs: Any) -> tuple[str, int]:
+        outputs = self.model.generate(**kwargs)[0]
+        completion = self.tokenizer.decode(outputs[prompt_token_count:], skip_special_tokens=True)
+
+        if kwargs["stopping_criteria"][0].__class__.__name__ == "StopSequenceCriteria":
+            for stop_sequence in kwargs["stopping_criteria"][0].stop_sequences:
+                completion = completion.split(stop_sequence)[0]
+        return completion, len(outputs[prompt_token_count:])
+
+    def logprobs(self, samples: list[Sample]) -> list[RawLoglikelihood]:
+        results = []
+        for sample in samples:
+            # format
+            prompt = self._formatter.format(sample.messages, output_mode="string")
+            choices_log_probs: dict[str, float] = {}
+            choices_log_probs_sequence_positions: dict[str, float] = {}
+            error: Error | None = None
+
+            for choice in sample.possible_completions or []:
+                num_choice_tokens = len(self.tokenizer.encode(choice, add_special_tokens=False))
+                prompt_and_choice = f"{prompt}{choice}"
+
+                total_tokens_count = len(self.tokenizer.encode(prompt_and_choice, add_special_tokens=False))
+
+                min_max_tokens = min(filter(None, [self.SEQ_LENGTH, self.seq_length]))
+
+                if min_max_tokens < total_tokens_count:
+                    if raise_errors():
+                        raise PromptTooLongException("Prompt exceeded context size.")
+                    choices_log_probs = {}
+                    choices_log_probs_sequence_positions = {}
+                    error = Error(
+                        error_class=PromptTooLongException.__name__,
+                        message="Prompt and choice exceeded context size.",
+                        traceback="",
+                    )
+                    break
+                else:
+                    # Calculate log-likelihoods for each token in the completion
+                    sum_log_probs = self._model_log_probs(prompt_and_choice, num_choice_tokens)
+
+                choices_log_probs.update({choice: sum_log_probs})
+                choices_log_probs_sequence_positions.update({choice: num_choice_tokens})
+
+            results.append(
+                RawLoglikelihood(
+                    prompt=prompt,
+                    prompt_sequence_positions=len(self.tokenizer.encode(prompt, add_special_tokens=False)),
+                    concat_compression=ConcatCompression.calculate(
+                        sample.messages, count_tokens=self.count_tokens, choices=sample.possible_completions
+                    ),
+                    loglikelihoods=choices_log_probs,
+                    loglikelihoods_sequence_positions=choices_log_probs_sequence_positions,
+                    raw_loglikelihood_error=error,
+                )
+            )
+        return results
+
+    def _model_log_probs(self, prompt: str, num_choice_tokens: int) -> float:
+        with torch.no_grad():
+            inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.device)
+            outputs = self.model(**inputs, labels=inputs["input_ids"])
+            logits = outputs.logits[:, :-1, :].squeeze(0)
+            target_ids = inputs["input_ids"][:, 1:].squeeze(0)
+
+            token_loglikelihoods = []
+            for i in range(0, len(target_ids)):
+                token_id = target_ids[i].item()
+                token = self.tokenizer.decode([token_id])
+                loglikelihood = torch.log_softmax(logits[i], dim=-1)[token_id].item()
+                token_loglikelihoods.append({token: loglikelihood})
+
+            return sum([list(log_prob.values())[0] for log_prob in token_loglikelihoods[-num_choice_tokens:]])
+
+    @property
+    def seq_length(self) -> int | None:
+        config = self.model.config
+        return config.max_position_embeddings if hasattr(config, "max_position_embeddings") else None
+
+
+class HFLLM_from_name(HFLLM):
+    """
+    A generic class to create HFLLM instances from a given model name.
+    """
+
+    def __init__(self, model_name: str | None = None, formatter: str = "Llama3Formatter", **kwargs: Any) -> None:
+        if model_name is None:
+            raise ValueError("model_name is required")
+
+        self.LLM_NAME = model_name
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(self.LLM_NAME)
+        self.model = AutoModelForCausalLM.from_pretrained(self.LLM_NAME, device_map="auto")
+
+        # Lazy formatter initialization - only create the one we need
+        selected_formatter = self._get_formatter(formatter, model_name)
+
+        print(f"{RED}[ Model initialized --------------------- {RESET}{self.LLM_NAME} {RED}]{RESET}")
+        print(f"{RED}[ Formatter: {formatter} ]{RESET}")
+        self._set_formatter(selected_formatter)
+
+    def _get_formatter(self, formatter: str, model_name: str) -> Any:
+        """Get formatter instance based on formatter name."""
+        if formatter == "Llama3Formatter":
+            return Llama3Formatter()
+        elif formatter == "MistralFormatter":
+            from eval_framework.llm.mistral import MagistralFormatter
+
+            return MagistralFormatter(model_name)
+        elif formatter == "ConcatFormatter":
+            return ConcatFormatter()
+        elif formatter == "HFFormatter":
+            return HFFormatter(model_name)
+        else:
+            supported = ["Llama3Formatter", "MistralFormatter", "ConcatFormatter", "HFFormatter"]
+            raise ValueError(f"Unsupported formatter: {formatter}. Supported formatters: {supported}")
+
+
+class Pythia410m(HFLLM):
+    LLM_NAME = "EleutherAI/pythia-410m"
+    DEFAULT_FORMATTER = ConcatFormatter
+
+
+class SmolLM135M(HFLLM):
+    LLM_NAME = "HuggingFaceTB/SmolLM-135M"
+    DEFAULT_FORMATTER = ConcatFormatter
+
+
+class Smollm135MInstruct(HFLLM):
+    LLM_NAME = "HuggingFaceTB/SmolLM-135M-Instruct"
+    DEFAULT_FORMATTER = partial(HFFormatter, LLM_NAME)
+
+
+class Qwen3_0_6B(HFLLM):
+    LLM_NAME = "Qwen/Qwen3-0.6B"
+    DEFAULT_FORMATTER = partial(HFFormatter, LLM_NAME, chat_template_kwargs={"enable_thinking": True})

eval_framework/llm/mistral.py

@@ -0,0 +1,73 @@
+from functools import partial
+from typing import Any, Literal, override
+
+from vllm import SamplingParams
+
+from eval_framework.llm.vllm import TokenizedContainer, VLLMModel, VLLMTokenizerAPI
+from template_formatting.formatter import BaseFormatter, Message
+from template_formatting.mistral_formatter import MagistralFormatter, MistralSerializer
+
+__all__ = [
+    "MistralAdapter",
+    "MistralVLLM",
+]
+
+
+class MistralAdapter(VLLMTokenizerAPI[list[Message]]):
+    def __init__(self, target_mdl: str) -> None:
+        self.serializer = MistralSerializer(llm_target=target_mdl)
+        self.tokenizer = self.serializer.get_tokenizer()
+
+    def encode_formatted_struct(self, struct: list[Message]) -> TokenizedContainer:
+        mistral_msg_lst = self.serializer.convert_from_aa(msg_lst=struct)
+        mistral_request = self.serializer.build_mistral_request(mistral_msg_lst=mistral_msg_lst)
+        mistral_tokenized_obj = self.tokenizer.encode_instruct(mistral_request)
+        return TokenizedContainer(tokens=mistral_tokenized_obj.tokens, text=mistral_tokenized_obj.text)
+
+    def encode_plain_text(self, text: str) -> TokenizedContainer:
+        choice_tokens = self.tokenizer.tokenizer.encode(text, False, False)
+        return TokenizedContainer(tokens=choice_tokens, text=text)
+
+
+class MistralVLLM(VLLMModel):
+    def __init__(
+        self,
+        formatter: BaseFormatter | None = None,
+        max_model_len: int | None = None,
+        tensor_parallel_size: int = 1,
+        gpu_memory_utilization: float = 0.9,
+        batch_size: int = 1,
+        checkpoint_path: str | None = None,
+        checkpoint_name: str | None = None,
+        sampling_params: SamplingParams | dict[str, Any] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        model_args = {"tokenizer_mode": "mistral", "config_format": "mistral", "load_format": "mistral"}
+        super().__init__(
+            formatter,
+            max_model_len,
+            tensor_parallel_size,
+            gpu_memory_utilization,
+            batch_size,
+            checkpoint_path,
+            checkpoint_name,
+            sampling_params,
+            **{**model_args, **kwargs},
+        )
+
+    @override
+    @property
+    def tokenizer(self) -> VLLMTokenizerAPI:
+        if self._tokenizer is None:
+            self._tokenizer = MistralAdapter(target_mdl=self.LLM_NAME)
+        return self._tokenizer
+
+    @property
+    def formatter_output_mode(self) -> Literal["string", "list"]:
+        """Determine the correct output mode for the formatter based on tokenizer type."""
+        return "list"
+
+
+class MagistralVLLM(MistralVLLM):
+    LLM_NAME = "mistralai/Magistral-Small-2506"
+    DEFAULT_FORMATTER = partial(MagistralFormatter, "mistralai/Magistral-Small-2506")

eval_framework/llm/models.py

@@ -0,0 +1,16 @@
+"""This is just a default model file with some small models for testing.
+
+Please define your own model file externally and pass it to the eval-framework entrypoint
+to use it.
+"""
+
+from eval_framework.utils.packaging import is_extra_installed
+
+if is_extra_installed(extra="transformers"):
+    from eval_framework.llm.huggingface import Pythia410m, SmolLM135M, Smollm135MInstruct, Qwen3_0_6B  # noqa F401
+
+if is_extra_installed("mistral"):
+    from eval_framework.llm.mistral import MagistralVLLM  # noqa F401
+
+if is_extra_installed("vllm"):
+    from eval_framework.llm.vllm import Qwen3_0_6B_VLLM, Qwen3_0_6B_VLLM_No_Thinking  # noqa F401

eval_framework/llm/openai.py

@@ -0,0 +1,205 @@
+import json
+import logging
+import os
+from collections.abc import Callable, Sequence
+from typing import Any
+
+import tiktoken  # OpenAI's official tokenizer library
+from openai import OpenAI
+
+from eval_framework.llm.base import BaseLLM
+from eval_framework.shared.types import ConcatCompression, RawCompletion, RawLoglikelihood
+from eval_framework.tasks.base import Sample
+from template_formatting.formatter import BaseFormatter, Message, Role
+
+logger = logging.getLogger(__name__)
+
+
+class OpenAIModel(BaseLLM):
+    DEFAULT_FORMATTER: Callable[[], BaseFormatter] | None = None
+
+    def __init__(
+        self,
+        model_name: str = "gpt-4o",
+        formatter: BaseFormatter | None = None,
+        temperature: float | None = None,
+        api_key: str | None = None,
+        organization: str | None = None,
+        base_url: str | None = None,
+    ) -> None:
+        """Initialize OpenAI API client.
+        Args:
+            model_name: Name of the OpenAI model to use (e.g., "gpt-4", "gpt-3.5-turbo")
+            formatter: Optional message formatter
+            temperature: Sampling temperature (0.0 to 2.0)
+            api_key: OpenAI API key (defaults to OPENAI_API_KEY env variable)
+            organization: Optional organization ID
+            base_url: Optional API base URL for Azure or other endpoints
+        """
+        self._model_name = model_name
+        logger.info(f"Using {model_name} as a judge")
+        self._formatter = formatter or self.DEFAULT_FORMATTER() if self.DEFAULT_FORMATTER is not None else None
+        self._temperature = temperature
+        # Initialize OpenAI client
+        self._client = OpenAI(
+            api_key=api_key or os.getenv("OPENAI_API_KEY", ""),
+            organization=organization,
+            base_url=base_url,
+        )
+
+        # Initialize tiktoken tokenizer for the model
+        self._encoding = tiktoken.encoding_for_model(self._model_name)
+
+    def _count_tokens(self, text: str) -> int:
+        """Helper method to count tokens using tiktoken."""
+        return len(self._encoding.encode(text))
+
+    def generate_from_messages(
+        self,
+        messages: list[Sequence[Message]],
+        stop_sequences: list[str] | None = None,
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+    ) -> list[RawCompletion]:
+        if temperature is None:
+            effective_temperature = 0.0  # Current default, TODO: refactor to use model's default
+            logger.info(
+                f"Using default temperature value: {effective_temperature} as no custom temperature value was provided"
+            )
+        else:
+            effective_temperature = temperature
+        """Generate completion from messages.
+        Args:
+            messages: Sequence of messages
+            stop_sequences: Optional list of stop sequences
+            max_tokens: Optional maximum number of tokens to generate
+        Returns:
+            Tuple of (prompt, completion)
+        """
+        results = []
+        for single_messages in messages:
+            if self._formatter is not None:
+                # Use formatter for text completion API
+                prompt = self._formatter.format(single_messages, output_mode="string")
+                response = self._client.completions.create(
+                    model=self._model_name,
+                    prompt=prompt,
+                    temperature=effective_temperature,
+                    max_tokens=max_tokens,
+                    stop=stop_sequences,
+                )
+
+                prompt_sequence_positions: int | None = self._count_tokens(prompt)
+                completion = response.choices[0].text
+                completion_sequence_positions = self._count_tokens(completion)
+
+                results.append(
+                    RawCompletion(
+                        prompt=prompt,
+                        prompt_sequence_positions=prompt_sequence_positions,
+                        concat_compression=ConcatCompression.calculate(
+                            single_messages, count_tokens=self._count_tokens, completion=completion
+                        ),
+                        completion=completion,
+                        completion_sequence_positions=completion_sequence_positions,
+                    )
+                )
+            else:
+                # Use chat completion API
+                from openai.types.chat import ChatCompletionAssistantMessageParam, ChatCompletionUserMessageParam
+
+                chat_messages = [
+                    (
+                        ChatCompletionUserMessageParam(role="user", content=m.content)
+                        if m.role is not None and m.role.value.lower() == "user"
+                        else ChatCompletionAssistantMessageParam(role="assistant", content=m.content)
+                    )
+                    for m in single_messages
+                ]
+
+                chat_response = self._client.chat.completions.create(
+                    model=self._model_name,
+                    messages=chat_messages,
+                    temperature=effective_temperature,
+                    max_tokens=max_tokens,
+                    stop=stop_sequences,
+                )
+
+                # Reconstruct the prompt (since OpenAI API does not return it)
+                prompt = "\n".join([f"{m['role']}: {m['content']}" for m in chat_messages])
+
+                prompt_sequence_positions = (
+                    chat_response.usage.prompt_tokens if chat_response.usage else None
+                )  # OpenAI API gives token count
+                completion = (
+                    chat_response.choices[0].message.content if chat_response.choices[0].message.content else ""
+                )
+                completion_sequence_positions = self._count_tokens(completion)
+
+                results.append(
+                    RawCompletion(
+                        prompt=prompt,
+                        prompt_sequence_positions=prompt_sequence_positions,
+                        concat_compression=ConcatCompression.calculate(
+                            single_messages, count_tokens=self._count_tokens, completion=completion
+                        ),
+                        completion=completion,
+                        completion_sequence_positions=completion_sequence_positions,
+                    )
+                )
+        return results
+
+    def logprobs(self, samples: list[Sample]) -> list[RawLoglikelihood]:
+        """Get log probabilities for possible completions.
+        Args:
+            samples: list of Sample containing possible completions
+        Returns:
+            list of Tuple of (prompt, dict of completion->logprob)
+        Raises:
+            NotImplementedError: Logprobs not yet implemented
+        """
+        raise NotImplementedError("Logprobs not yet implemented for OpenAI API")
+
+    def generate_structured_output(
+        self,
+        messages: list[Sequence[Message]],
+        stop_sequences: list[str] | None = None,
+        max_tokens: int | None = None,
+        temperature: float = 0.0,
+    ) -> Any:
+        """Generate structured output (e.g. JSON) from messages.
+        This implementation ensures the model returns valid JSON.
+        Args:
+            messages: list of Sequence of messages
+            stop_sequences: Optional stop sequences
+            max_tokens: Optional max tokens
+        Returns:
+            Parsed JSON response
+        """
+        completions = []
+        list_json_messages: list[Sequence[Message]] = []
+        for single_messages in messages:
+            # Add system message to encourage JSON output
+            json_messages = list(single_messages)
+            if not any(m.role == Role.SYSTEM for m in single_messages):
+                json_messages.insert(
+                    0,
+                    Message(
+                        role=Role.SYSTEM, content="You are a helpful assistant that always responds with valid JSON."
+                    ),
+                )
+            list_json_messages.append(json_messages)
+        # Generate completion
+        completions = self.generate_from_messages(
+            messages=list_json_messages, stop_sequences=stop_sequences, max_tokens=max_tokens
+        )
+        responses = []
+        for completion in completions:
+            try:
+                # Parse JSON responses
+                responses.append(json.loads(completion.completion))
+            except json.JSONDecodeError as e:
+                logger.info(f"Warning: Failed to parse JSON response: {e}")
+                logger.info(f"Raw response: {completion.completion}")
+                raise
+        return responses