diff --git a/README.md b/README.md
index 42aa8a1..eb0e61f 100644
--- a/README.md
+++ b/README.md
@@ -36,6 +36,22 @@ python jsonmode.py --query "Please return a json object to represent Goku from t
 - `--load_in_4bit`: Option to load in 4bit with bitsandbytes (default: "False").
 - `--query`: Query to be used for function call inference (default: "I need the current stock price of Tesla (TSLA)").
 - `--max_depth`: Maximum number of recursive iterations (default: 5).
+- `--attn_implementation`: Attention implementation passed to Transformers (default: "flash_attention_2"). Use "eager" or "sdpa" if FlashAttention is not compatible with your PyTorch/CUDA environment.
+
+### FlashAttention startup errors
+
+The example scripts default to `flash_attention_2` for faster inference. If startup fails with an error from `flash_attn` or `flash_attn_2_cuda`, the installed FlashAttention wheel usually does not match the active PyTorch/CUDA build. Reinstall FlashAttention in the current environment:
+
+```bash
+pip uninstall flash-attn
+pip install --no-build-isolation flash-attn
+```
+
+If you want to run without FlashAttention, disable it explicitly:
+
+```bash
+python functioncall.py --attn_implementation eager --query "I need the current stock price of Tesla (TSLA)"
+```
 
 ## Adding Custom Functions
 
diff --git a/functioncall.py b/functioncall.py
index 7d0d544..dfbf2f0 100644
--- a/functioncall.py
+++ b/functioncall.py
@@ -1,91 +1,143 @@
 import argparse
-import torch
 import json
 
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig
+FLASH_ATTENTION_TROUBLESHOOTING = (
+    "Failed to load the model with FlashAttention 2. This usually means the installed "
+    "flash-attn wheel is not ABI-compatible with your PyTorch/CUDA build. Reinstall "
+    "flash-attn for the active environment, for example `pip uninstall flash-attn && "
+    "pip install --no-build-isolation flash-attn`, or run with `--attn_implementation "
+    "eager` to disable FlashAttention."
 )
 
-import functions
-from prompter import PromptManager
-from validator import validate_function_call_schema
 
-from utils import (
-    print_nous_text_art,
-    inference_logger,
-    get_assistant_message,
-    get_chat_template,
-    validate_and_extract_tool_calls
-)
+def is_flash_attention_import_error(error):
+    message = str(error).lower()
+    return (
+        "flash_attn" in message
+        or "flash-attn" in message
+        or "flash attention" in message
+    )
+
+
+def load_runtime_dependencies():
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+    import functions
+    from prompter import PromptManager
+    from validator import validate_function_call_schema
+    from utils import (
+        print_nous_text_art,
+        inference_logger,
+        get_assistant_message,
+        get_chat_template,
+        validate_and_extract_tool_calls,
+    )
+
+    return {
+        "torch": torch,
+        "AutoModelForCausalLM": AutoModelForCausalLM,
+        "AutoTokenizer": AutoTokenizer,
+        "BitsAndBytesConfig": BitsAndBytesConfig,
+        "functions": functions,
+        "PromptManager": PromptManager,
+        "validate_function_call_schema": validate_function_call_schema,
+        "print_nous_text_art": print_nous_text_art,
+        "inference_logger": inference_logger,
+        "get_assistant_message": get_assistant_message,
+        "get_chat_template": get_chat_template,
+        "validate_and_extract_tool_calls": validate_and_extract_tool_calls,
+    }
+
 
 class ModelInference:
-    def __init__(self, model_path, chat_template, load_in_4bit):
-        inference_logger.info(print_nous_text_art())
-        self.prompter = PromptManager()
+    def __init__(self, model_path, chat_template, load_in_4bit, attn_implementation):
+        deps = load_runtime_dependencies()
+        self.functions = deps["functions"]
+        self.validate_function_call_schema = deps["validate_function_call_schema"]
+        self.inference_logger = deps["inference_logger"]
+        self.get_assistant_message = deps["get_assistant_message"]
+        self.get_chat_template = deps["get_chat_template"]
+        self.validate_and_extract_tool_calls = deps["validate_and_extract_tool_calls"]
+
+        self.inference_logger.info(deps["print_nous_text_art"]())
+        self.prompter = deps["PromptManager"]()
         self.bnb_config = None
 
         if load_in_4bit == "True":
-            self.bnb_config = BitsAndBytesConfig(
+            self.bnb_config = deps["BitsAndBytesConfig"](
                 load_in_4bit=True,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_use_double_quant=True,
             )
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            trust_remote_code=True,
-            return_dict=True,
-            quantization_config=self.bnb_config,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2",
-            device_map="auto",
-        )
+        try:
+            self.model = deps["AutoModelForCausalLM"].from_pretrained(
+                model_path,
+                trust_remote_code=True,
+                return_dict=True,
+                quantization_config=self.bnb_config,
+                torch_dtype=deps["torch"].float16,
+                attn_implementation=attn_implementation,
+                device_map="auto",
+            )
+        except (ImportError, OSError, RuntimeError) as exc:
+            if (
+                attn_implementation == "flash_attention_2"
+                and is_flash_attention_import_error(exc)
+            ):
+                raise RuntimeError(FLASH_ATTENTION_TROUBLESHOOTING) from exc
+            raise
 
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+        self.tokenizer = deps["AutoTokenizer"].from_pretrained(
+            model_path, trust_remote_code=True
+        )
         self.tokenizer.pad_token = self.tokenizer.eos_token
         self.tokenizer.padding_side = "left"
 
         if self.tokenizer.chat_template is None:
             print("No chat template defined, getting chat_template...")
-            self.tokenizer.chat_template = get_chat_template(chat_template)
-        
-        inference_logger.info(self.model.config)
-        inference_logger.info(self.model.generation_config)
-        inference_logger.info(self.tokenizer.special_tokens_map)
+            self.tokenizer.chat_template = self.get_chat_template(chat_template)
+
+        self.inference_logger.info(self.model.config)
+        self.inference_logger.info(self.model.generation_config)
+        self.inference_logger.info(self.tokenizer.special_tokens_map)
 
     def process_completion_and_validate(self, completion, chat_template):
 
-        assistant_message = get_assistant_message(completion, chat_template, self.tokenizer.eos_token)
+        assistant_message = self.get_assistant_message(
+            completion, chat_template, self.tokenizer.eos_token
+        )
 
         if assistant_message:
-            validation, tool_calls, error_message = validate_and_extract_tool_calls(assistant_message)
+            validation, tool_calls, error_message = (
+                self.validate_and_extract_tool_calls(assistant_message)
+            )
 
             if validation:
-                inference_logger.info(f"parsed tool calls:\n{json.dumps(tool_calls, indent=2)}")
+                self.inference_logger.info(
+                    f"parsed tool calls:\n{json.dumps(tool_calls, indent=2)}"
+                )
                 return tool_calls, assistant_message, error_message
             else:
                 tool_calls = None
                 return tool_calls, assistant_message, error_message
         else:
-            inference_logger.warning("Assistant message is None")
+            self.inference_logger.warning("Assistant message is None")
             raise ValueError("Assistant message is None")
-        
+
     def execute_function_call(self, tool_call):
         function_name = tool_call.get("name")
-        function_to_call = getattr(functions, function_name, None)
+        function_to_call = getattr(self.functions, function_name, None)
         function_args = tool_call.get("arguments", {})
 
-        inference_logger.info(f"Invoking function call {function_name} ...")
+        self.inference_logger.info(f"Invoking function call {function_name} ...")
         function_response = function_to_call(*function_args.values())
         results_dict = f'{{"name": "{function_name}", "content": {function_response}}}'
         return results_dict
-    
+
     def run_inference(self, prompt):
         inputs = self.tokenizer.apply_chat_template(
-            prompt,
-            add_generation_prompt=True,
-            return_tensors='pt'
+            prompt, add_generation_prompt=True, return_tensors="pt"
         )
 
         tokens = self.model.generate(
@@ -94,9 +146,11 @@ def run_inference(self, prompt):
             temperature=0.8,
             repetition_penalty=1.1,
             do_sample=True,
-            eos_token_id=self.tokenizer.eos_token_id
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        completion = self.tokenizer.decode(
+            tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True
         )
-        completion = self.tokenizer.decode(tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True)
         return completion
 
     def generate_function_call(self, query, chat_template, num_fewshot, max_depth=5):
@@ -104,78 +158,133 @@ def generate_function_call(self, query, chat_template, num_fewshot, max_depth=5)
             depth = 0
             user_message = f"{query}\nThis is the first turn and you don't have <tool_results> to analyze yet"
             chat = [{"role": "user", "content": user_message}]
-            tools = functions.get_openai_tools()
+            tools = self.functions.get_openai_tools()
             prompt = self.prompter.generate_prompt(chat, tools, num_fewshot)
             completion = self.run_inference(prompt)
 
             def recursive_loop(prompt, completion, depth):
                 nonlocal max_depth
-                tool_calls, assistant_message, error_message = self.process_completion_and_validate(completion, chat_template)
+                tool_calls, assistant_message, error_message = (
+                    self.process_completion_and_validate(completion, chat_template)
+                )
                 prompt.append({"role": "assistant", "content": assistant_message})
 
-                tool_message = f"Agent iteration {depth} to assist with user query: {query}\n"
+                tool_message = (
+                    f"Agent iteration {depth} to assist with user query: {query}\n"
+                )
                 if tool_calls:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    self.inference_logger.info(
+                        f"Assistant Message:\n{assistant_message}"
+                    )
 
                     for tool_call in tool_calls:
-                        validation, message = validate_function_call_schema(tool_call, tools)
+                        validation, message = self.validate_function_call_schema(
+                            tool_call, tools
+                        )
                         if validation:
                             try:
-                                function_response = self.execute_function_call(tool_call)
+                                function_response = self.execute_function_call(
+                                    tool_call
+                                )
                                 tool_message += f"<tool_response>\n{function_response}\n</tool_response>\n"
-                                inference_logger.info(f"Here's the response from the function call: {tool_call.get('name')}\n{function_response}")
+                                self.inference_logger.info(
+                                    f"Here's the response from the function call: {tool_call.get('name')}\n{function_response}"
+                                )
                             except Exception as e:
-                                inference_logger.info(f"Could not execute function: {e}")
+                                self.inference_logger.info(
+                                    f"Could not execute function: {e}"
+                                )
                                 tool_message += f"<tool_response>\nThere was an error when executing the function: {tool_call.get('name')}\nHere's the error traceback: {e}\nPlease call this function again with correct arguments within XML tags <tool_call></tool_call>\n</tool_response>\n"
                         else:
-                            inference_logger.info(message)
+                            self.inference_logger.info(message)
                             tool_message += f"<tool_response>\nThere was an error validating function call against function signature: {tool_call.get('name')}\nHere's the error traceback: {message}\nPlease call this function again with correct arguments within XML tags <tool_call></tool_call>\n</tool_response>\n"
                     prompt.append({"role": "tool", "content": tool_message})
 
                     depth += 1
                     if depth >= max_depth:
-                        print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
+                        print(
+                            f"Maximum recursion depth reached ({max_depth}). Stopping recursion."
+                        )
                         return
 
                     completion = self.run_inference(prompt)
                     recursive_loop(prompt, completion, depth)
                 elif error_message:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    self.inference_logger.info(
+                        f"Assistant Message:\n{assistant_message}"
+                    )
                     tool_message += f"<tool_response>\nThere was an error parsing function calls\n Here's the error stack trace: {error_message}\nPlease call the function again with correct syntax<tool_response>"
                     prompt.append({"role": "tool", "content": tool_message})
 
                     depth += 1
                     if depth >= max_depth:
-                        print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
+                        print(
+                            f"Maximum recursion depth reached ({max_depth}). Stopping recursion."
+                        )
                         return
 
                     completion = self.run_inference(prompt)
                     recursive_loop(prompt, completion, depth)
                 else:
-                    inference_logger.info(f"Assistant Message:\n{assistant_message}")
+                    self.inference_logger.info(
+                        f"Assistant Message:\n{assistant_message}"
+                    )
 
             recursive_loop(prompt, completion, depth)
 
         except Exception as e:
-            inference_logger.error(f"Exception occurred: {e}")
+            self.inference_logger.error(f"Exception occurred: {e}")
             raise e
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run recursive function calling loop")
     parser.add_argument("--model_path", type=str, help="Path to the model folder")
-    parser.add_argument("--chat_template", type=str, default="chatml", help="Chat template for prompt formatting")
-    parser.add_argument("--num_fewshot", type=int, default=None, help="Option to use json mode examples")
-    parser.add_argument("--load_in_4bit", type=str, default="False", help="Option to load in 4bit with bitsandbytes")
-    parser.add_argument("--query", type=str, default="I need the current stock price of Tesla (TSLA)")
-    parser.add_argument("--max_depth", type=int, default=5, help="Maximum number of recursive iteration")
+    parser.add_argument(
+        "--chat_template",
+        type=str,
+        default="chatml",
+        help="Chat template for prompt formatting",
+    )
+    parser.add_argument(
+        "--num_fewshot", type=int, default=None, help="Option to use json mode examples"
+    )
+    parser.add_argument(
+        "--load_in_4bit",
+        type=str,
+        default="False",
+        help="Option to load in 4bit with bitsandbytes",
+    )
+    parser.add_argument(
+        "--query", type=str, default="I need the current stock price of Tesla (TSLA)"
+    )
+    parser.add_argument(
+        "--max_depth", type=int, default=5, help="Maximum number of recursive iteration"
+    )
+    parser.add_argument(
+        "--attn_implementation",
+        type=str,
+        default="flash_attention_2",
+        choices=["flash_attention_2", "eager", "sdpa"],
+        help="Attention implementation passed to transformers.from_pretrained",
+    )
     args = parser.parse_args()
 
     # specify custom model path
     if args.model_path:
-        inference = ModelInference(args.model_path, args.chat_template, args.load_in_4bit)
+        inference = ModelInference(
+            args.model_path,
+            args.chat_template,
+            args.load_in_4bit,
+            args.attn_implementation,
+        )
     else:
-        model_path = 'NousResearch/Hermes-2-Pro-Llama-3-8B'
-        inference = ModelInference(model_path, args.chat_template, args.load_in_4bit)
-        
+        model_path = "NousResearch/Hermes-2-Pro-Llama-3-8B"
+        inference = ModelInference(
+            model_path, args.chat_template, args.load_in_4bit, args.attn_implementation
+        )
+
     # Run the model evaluator
-    inference.generate_function_call(args.query, args.chat_template, args.num_fewshot, args.max_depth)
+    inference.generate_function_call(
+        args.query, args.chat_template, args.num_fewshot, args.max_depth
+    )
diff --git a/jsonmode.py b/jsonmode.py
index 4a569e5..1ecff1a 100644
--- a/jsonmode.py
+++ b/jsonmode.py
@@ -1,80 +1,119 @@
 import argparse
-import torch
 import json
 
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    BitsAndBytesConfig
+FLASH_ATTENTION_TROUBLESHOOTING = (
+    "Failed to load the model with FlashAttention 2. This usually means the installed "
+    "flash-attn wheel is not ABI-compatible with your PyTorch/CUDA build. Reinstall "
+    "flash-attn for the active environment, for example `pip uninstall flash-attn && "
+    "pip install --no-build-isolation flash-attn`, or run with `--attn_implementation "
+    "eager` to disable FlashAttention."
 )
 
-from validator import validate_json_data
 
-from utils import (
-    print_nous_text_art,
-    inference_logger,
-    get_assistant_message,
-    get_chat_template,
-    validate_and_extract_tool_calls
-)
+def is_flash_attention_import_error(error):
+    message = str(error).lower()
+    return (
+        "flash_attn" in message
+        or "flash-attn" in message
+        or "flash attention" in message
+    )
+
+
+def load_runtime_dependencies():
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+    from validator import validate_json_data
+    from utils import (
+        print_nous_text_art,
+        inference_logger,
+        get_assistant_message,
+        get_chat_template,
+    )
+
+    return {
+        "torch": torch,
+        "AutoModelForCausalLM": AutoModelForCausalLM,
+        "AutoTokenizer": AutoTokenizer,
+        "BitsAndBytesConfig": BitsAndBytesConfig,
+        "validate_json_data": validate_json_data,
+        "print_nous_text_art": print_nous_text_art,
+        "inference_logger": inference_logger,
+        "get_assistant_message": get_assistant_message,
+        "get_chat_template": get_chat_template,
+    }
+
 
-# create your pydantic model for json object here
-from typing import List, Optional
-from pydantic import BaseModel
+def get_pydantic_schema():
+    # Create your pydantic model for json object here.
+    from typing import List, Optional
+    from pydantic import BaseModel, ConfigDict
 
-class Character(BaseModel):
-    name: str
-    species: str
-    role: str
-    personality_traits: Optional[List[str]]
-    special_attacks: Optional[List[str]]
+    class Character(BaseModel):
+        name: str
+        species: str
+        role: str
+        personality_traits: Optional[List[str]]
+        special_attacks: Optional[List[str]]
 
-    class Config:
-        schema_extra = {
-            "additionalProperties": False
-        }
+        model_config = ConfigDict(json_schema_extra={"additionalProperties": False})
+
+    # Serialize the pydantic model into json schema.
+    return json.dumps(Character.model_json_schema())
 
-# serialize pydantic model into json schema
-pydantic_schema = Character.schema_json()
 
 class ModelInference:
-    def __init__(self, model_path, chat_template, load_in_4bit):
-        inference_logger.info(print_nous_text_art())
+    def __init__(self, model_path, chat_template, load_in_4bit, attn_implementation):
+        deps = load_runtime_dependencies()
+        self.validate_json_data = deps["validate_json_data"]
+        self.inference_logger = deps["inference_logger"]
+        self.get_assistant_message = deps["get_assistant_message"]
+        self.get_chat_template = deps["get_chat_template"]
+
+        self.inference_logger.info(deps["print_nous_text_art"]())
         self.bnb_config = None
 
         if load_in_4bit == "True":
-            self.bnb_config = BitsAndBytesConfig(
+            self.bnb_config = deps["BitsAndBytesConfig"](
                 load_in_4bit=True,
                 bnb_4bit_quant_type="nf4",
                 bnb_4bit_use_double_quant=True,
             )
-        self.model = AutoModelForCausalLM.from_pretrained(
-            model_path,
-            trust_remote_code=True,
-            return_dict=True,
-            quantization_config=self.bnb_config,
-            torch_dtype=torch.float16,
-            attn_implementation="flash_attention_2",
-            device_map="auto",
+        try:
+            self.model = deps["AutoModelForCausalLM"].from_pretrained(
+                model_path,
+                trust_remote_code=True,
+                return_dict=True,
+                quantization_config=self.bnb_config,
+                torch_dtype=deps["torch"].float16,
+                attn_implementation=attn_implementation,
+                device_map="auto",
+            )
+        except (ImportError, OSError, RuntimeError) as exc:
+            if (
+                attn_implementation == "flash_attention_2"
+                and is_flash_attention_import_error(exc)
+            ):
+                raise RuntimeError(FLASH_ATTENTION_TROUBLESHOOTING) from exc
+            raise
+
+        self.tokenizer = deps["AutoTokenizer"].from_pretrained(
+            model_path, trust_remote_code=True
         )
-
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
         self.tokenizer.pad_token = self.tokenizer.eos_token
         self.tokenizer.padding_side = "left"
 
         if self.tokenizer.chat_template is None:
             print("No chat template defined, getting chat_template...")
-            self.tokenizer.chat_template = get_chat_template(chat_template)
-        
-        inference_logger.info(self.model.config)
-        inference_logger.info(self.model.generation_config)
-        inference_logger.info(self.tokenizer.special_tokens_map)
-    
+            self.tokenizer.chat_template = self.get_chat_template(chat_template)
+
+        self.inference_logger.info(self.model.config)
+        self.inference_logger.info(self.model.generation_config)
+        self.inference_logger.info(self.tokenizer.special_tokens_map)
+
     def run_inference(self, prompt):
         inputs = self.tokenizer.apply_chat_template(
-            prompt,
-            add_generation_prompt=True,
-            return_tensors='pt'
+            prompt, add_generation_prompt=True, return_tensors="pt"
         )
 
         tokens = self.model.generate(
@@ -83,68 +122,119 @@ def run_inference(self, prompt):
             temperature=0.8,
             repetition_penalty=1.1,
             do_sample=True,
-            eos_token_id=self.tokenizer.eos_token_id
+            eos_token_id=self.tokenizer.eos_token_id,
+        )
+        completion = self.tokenizer.decode(
+            tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True
         )
-        completion = self.tokenizer.decode(tokens[0], skip_special_tokens=False, clean_up_tokenization_space=True)
         return completion
 
     def generate_json_completion(self, query, chat_template, max_depth=5):
         try:
+            pydantic_schema = get_pydantic_schema()
             depth = 0
             sys_prompt = f"You are a helpful assistant that answers in JSON. Here's the json schema you must adhere to:\n<schema>\n{pydantic_schema}\n</schema>"
             prompt = [{"role": "system", "content": sys_prompt}]
             prompt.append({"role": "user", "content": query})
 
-            inference_logger.info(f"Running inference to generate json object for pydantic schema:\n{json.dumps(json.loads(pydantic_schema), indent=2)}")
+            self.inference_logger.info(
+                f"Running inference to generate json object for pydantic schema:\n{json.dumps(json.loads(pydantic_schema), indent=2)}"
+            )
             completion = self.run_inference(prompt)
 
             def recursive_loop(prompt, completion, depth):
                 nonlocal max_depth
 
-                assistant_message = get_assistant_message(completion, chat_template, self.tokenizer.eos_token)
+                assistant_message = self.get_assistant_message(
+                    completion, chat_template, self.tokenizer.eos_token
+                )
 
-                tool_message = f"Agent iteration {depth} to assist with user query: {query}\n"
+                tool_message = (
+                    f"Agent iteration {depth} to assist with user query: {query}\n"
+                )
                 if assistant_message is not None:
-                    validation, json_object, error_message = validate_json_data(assistant_message, json.loads(pydantic_schema))
+                    validation, json_object, error_message = self.validate_json_data(
+                        assistant_message, json.loads(pydantic_schema)
+                    )
                     if validation:
-                        inference_logger.info(f"Assistant Message:\n{assistant_message}")
-                        inference_logger.info(f"json schema validation passed")
-                        inference_logger.info(f"parsed json object:\n{json.dumps(json_object, indent=2)}")
+                        self.inference_logger.info(
+                            f"Assistant Message:\n{assistant_message}"
+                        )
+                        self.inference_logger.info("json schema validation passed")
+                        self.inference_logger.info(
+                            f"parsed json object:\n{json.dumps(json_object, indent=2)}"
+                        )
                     elif error_message:
-                        inference_logger.info(f"Assistant Message:\n{assistant_message}")
-                        inference_logger.info(f"json schema validation failed")
+                        self.inference_logger.info(
+                            f"Assistant Message:\n{assistant_message}"
+                        )
+                        self.inference_logger.info("json schema validation failed")
                         tool_message += f"<tool_response>\nJson schema validation failed\nHere's the error stacktrace: {error_message}\nPlease return corrrect json object\n<tool_response>"
-                        
+
                         depth += 1
                         if depth >= max_depth:
-                            print(f"Maximum recursion depth reached ({max_depth}). Stopping recursion.")
+                            print(
+                                f"Maximum recursion depth reached ({max_depth}). Stopping recursion."
+                            )
                             return
-                        
+
                         prompt.append({"role": "tool", "content": tool_message})
                         completion = self.run_inference(prompt)
                         recursive_loop(prompt, completion, depth)
                 else:
-                    inference_logger.warning("Assistant message is None")
+                    self.inference_logger.warning("Assistant message is None")
+
             recursive_loop(prompt, completion, depth)
         except Exception as e:
-            inference_logger.error(f"Exception occurred: {e}")
+            self.inference_logger.error(f"Exception occurred: {e}")
             raise e
 
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Run json mode completion")
     parser.add_argument("--model_path", type=str, help="Path to the model folder")
-    parser.add_argument("--chat_template", type=str, default="chatml", help="Chat template for prompt formatting")
-    parser.add_argument("--load_in_4bit", type=str, default="False", help="Option to load in 4bit with bitsandbytes")
-    parser.add_argument("--query", type=str, default="Please return a json object to represent Goku from the anime Dragon Ball Z?")
-    parser.add_argument("--max_depth", type=int, default=5, help="Maximum number of recursive iteration")
+    parser.add_argument(
+        "--chat_template",
+        type=str,
+        default="chatml",
+        help="Chat template for prompt formatting",
+    )
+    parser.add_argument(
+        "--load_in_4bit",
+        type=str,
+        default="False",
+        help="Option to load in 4bit with bitsandbytes",
+    )
+    parser.add_argument(
+        "--query",
+        type=str,
+        default="Please return a json object to represent Goku from the anime Dragon Ball Z?",
+    )
+    parser.add_argument(
+        "--max_depth", type=int, default=5, help="Maximum number of recursive iteration"
+    )
+    parser.add_argument(
+        "--attn_implementation",
+        type=str,
+        default="flash_attention_2",
+        choices=["flash_attention_2", "eager", "sdpa"],
+        help="Attention implementation passed to transformers.from_pretrained",
+    )
     args = parser.parse_args()
 
     # specify custom model path
     if args.model_path:
-        inference = ModelInference(args.model_path, args.chat_template, args.load_in_4bit)
+        inference = ModelInference(
+            args.model_path,
+            args.chat_template,
+            args.load_in_4bit,
+            args.attn_implementation,
+        )
     else:
-        model_path = 'NousResearch/Hermes-2-Pro-Llama-3-8B'
-        inference = ModelInference(model_path, args.chat_template, args.load_in_4bit)
-        
+        model_path = "NousResearch/Hermes-2-Pro-Llama-3-8B"
+        inference = ModelInference(
+            model_path, args.chat_template, args.load_in_4bit, args.attn_implementation
+        )
+
     # Run the model evaluator
     inference.generate_json_completion(args.query, args.chat_template, args.max_depth)