Spaces:
Runtime error
Runtime error
File size: 3,744 Bytes
200cf42 0cc0e0f 200cf42 7423671 200cf42 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import logging
from pathlib import Path
from typing import List, Tuple
import fire
logger = logging.getLogger(__name__)
# === Bot Loaders ===
def load_bot(
# env_file_path: str = ".env",
logging_config_path: str = "logging.yaml",
model_cache_dir: str = "./model_cache",
embedding_model_device: str = "cuda:0",
debug: bool = False,
):
"""
Load the financial assistant bot in production or development mode based on the `debug` flag
In DEV mode the embedding model runs on CPU and the fine-tuned LLM is mocked.
Otherwise, the embedding model runs on GPU and the fine-tuned LLM is used.
Args:
env_file_path (str): Path to the environment file.
logging_config_path (str): Path to the logging configuration file.
model_cache_dir (str): Path to the directory where the model cache is stored.
embedding_model_device (str): Device to use for the embedding model.
debug (bool): Flag to indicate whether to run the bot in debug mode or not.
Returns:
FinancialBot: An instance of the FinancialBot class.
"""
from financial_bot import initialize
# Be sure to initialize the environment variables before importing any other modules.
initialize(logging_config_path=logging_config_path, env_file_path=env_file_path)
from financial_bot import utils
from financial_bot.langchain_bot import FinancialBot
logger.info("#" * 100)
utils.log_available_gpu_memory()
utils.log_available_ram()
logger.info("#" * 100)
bot = FinancialBot(
model_cache_dir=Path(model_cache_dir) if model_cache_dir else None,
embedding_model_device=embedding_model_device,
debug=debug,
)
return bot
# === Bot Runners ===
def run_local(
about_me: str,
question: str,
history: List[Tuple[str, str]] = None,
debug: bool = False,
):
"""
Run the bot locally in production or dev mode.
Args:
about_me (str): A string containing information about the user.
question (str): A string containing the user's question.
history (List[Tuple[str, str]], optional): A list of tuples containing the user's previous questions
and the bot's responses. Defaults to None.
debug (bool, optional): A boolean indicating whether to run the bot in debug mode. Defaults to False.
Returns:
str: A string containing the bot's response to the user's question.
"""
if debug is True:
bot = load_bot_dev(model_cache_dir=None)
else:
bot = load_bot(model_cache_dir=None)
inputs = {
"about_me": about_me,
"question": question,
"history": history,
"context": bot,
}
response = _run(**inputs)
return response
def _run(**inputs):
"""
Central function that calls the bot and returns the response.
Args:
inputs (dict): A dictionary containing the following keys:
- context: The bot instance.
- about_me (str): Information about the user.
- question (str): The user's question.
- history (list): A list of previous conversations (optional).
Returns:
str: The bot's response to the user's question.
"""
from financial_bot import utils
logger.info("#" * 100)
utils.log_available_gpu_memory()
utils.log_available_ram()
logger.info("#" * 100)
bot = inputs["context"]
input_payload = {
"about_me": inputs["about_me"],
"question": inputs["question"],
"to_load_history": inputs["history"] if "history" in inputs else [],
}
response = bot.answer(**input_payload)
return response
if __name__ == "__main__":
fire.Fire(run_local)
|