functions/llm_chain.py

import transformers
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, AutoConfig, AutoModel
from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain.chains.llm import LLMChain
from langchain.memory import ConversationBufferWindowMemory
import torch
import datetime
from typing import Any, Dict, Union
from functions.context_engineering import get_context_data
import os
from safetensors.torch import load_model, save_model

def load_model(model_id: str = "teknium/OpenHermes-2.5-Mistral-7B") -> tuple:
    """
    Load the LLM and its corresponding tokenizer.

    Args:
        model_id (str, optional): Identifier for the pre-trained model. Defaults to "teknium/OpenHermes-2.5-Mistral-7B".

    Returns:
        tuple: A tuple containing the loaded model and tokenizer.
    """

    # Load the tokenizer for Mistral-7B-Instruct model
    tokenizer_path = "./mistral/tokenizer"
    if os.path.isdir(tokenizer_path) == False:
        tokenizer = AutoTokenizer.from_pretrained(model_id)
        tokenizer.save_pretrained(tokenizer_path)
    else:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)

    # Set the pad token to the unknown token to handle padding
    tokenizer.pad_token = tokenizer.unk_token

    # Set the padding side to "right" to prevent warnings during tokenization
    tokenizer.padding_side = "right"

    # BitsAndBytesConfig int-4 config
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    model_path = "/tmp/mistral/model"
    if os.path.exists(model_path):
        print("Loading model from disk")
        model_llm = AutoModelForCausalLM.from_pretrained(model_path)
    else:
        # Load the Mistral-7B-Instruct model with quantization configuration
        model_llm = AutoModelForCausalLM.from_pretrained(
            model_id,
            device_map="auto",
            quantization_config=bnb_config,
        )
        model_llm.save_pretrained(model_path)


    # Configure the pad token ID in the model to match the tokenizer's pad token ID
    model_llm.config.pad_token_id = tokenizer.pad_token_id

    return model_llm, tokenizer


def get_prompt_template():
    """
    Retrieve a template for generating prompts in a conversational AI system.

    Returns:
        str: A string representing the template for generating prompts.
            This template includes placeholders for system information,
            instructions, previous conversation, context, date and user query.
    """
    prompt_template = """<|im_start|>system
You are one of the best air quality experts in the world.

###INSTRUCTIONS:
- If you don't know the answer, you will respond politely that you cannot help.
- Use the context table with air quality indicators for city provided by user to generate your answer.
- You answer should be at least one sentence.
- Do not show any calculations to the user.
- Make sure that you use correct air quality indicators for the corresponding date.
- Add a rich analysis of the air quality level, such as whether it is safe, whether to go for a walk, etc.
- Do not mention in your answer that you are using context table.
<|im_end|>

### CONTEXT:
{context}

IMPORTANT: Today is {date_today}.

<|im_start|>user
{question}<|im_end|>
<|im_start|>assistant"""
    return prompt_template


def get_llm_chain(model_llm, tokenizer):
    """
    Create and configure a language model chain.

    Args:
        model_llm: The pre-trained language model for text generation.
        tokenizer: The tokenizer corresponding to the language model.

    Returns:
        LLMChain: The configured language model chain.
    """
    # Create a text generation pipeline using the loaded model and tokenizer
    text_generation_pipeline = transformers.pipeline(
        model=model_llm,                      # The pre-trained language model for text generation
        tokenizer=tokenizer,                  # The tokenizer corresponding to the language model
        task="text-generation",               # Specify the task as text generation
        use_cache=True,
        do_sample=True,
        temperature=0.4,
        top_p=1.0,
        top_k=0,
        max_new_tokens=512,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )

    # Create a Hugging Face pipeline for Mistral LLM using the text generation pipeline
    mistral_llm = HuggingFacePipeline(
        pipeline=text_generation_pipeline,
    )

    # Create prompt from prompt template
    prompt = PromptTemplate(
        input_variables=["context", "question", "date_today"],
        template=get_prompt_template(),
    )

    # Create LLM chain
    llm_chain = LLMChain(
        llm=mistral_llm,
        prompt=prompt,
        verbose=False,
    )

    return llm_chain


def generate_response(
    user_query: str,
    feature_view,
    weather_fg,
    model_air_quality,
    model_llm,
    tokenizer,
    llm_chain=None,
    verbose: bool = False,
) -> str:
    """
    Generate response to user query using LLM chain and context data.

    Args:
        user_query (str): The user's query.
        feature_view: Feature view for data retrieval.
        model_llm: Language model for text generation.
        tokenizer: Tokenizer for processing text.
        model_air_quality: Model for predicting air quality.
        llm_chain: LLM Chain.
        verbose (bool): Whether to print verbose information. Defaults to False.

    Returns:
        str: Generated response to the user query.
    """
    # Get context data based on user query
    context = get_context_data(
        user_query,
        feature_view,
        weather_fg,
        model_air_quality,
        model_llm=model_llm,
        tokenizer=tokenizer,
    )

    # Get today's date in a readable format
    date_today = f'{datetime.date.today().strftime("%A")}, {datetime.date.today()}'

    # Print today's date and context information if verbose mode is enabled
    if verbose:
        print(f"🗓️ Today's date: {date_today}")
        print(f'📖 {context}')

    # Invoke the language model chain with relevant context
    model_output = llm_chain.invoke({
        "context": context,
        "date_today": date_today,
        "question": user_query,
    })

    # Return the generated text from the model output
    return model_output['text'].split('<|im_start|>assistant')[-1]


def generate_response_openai(
    user_query: str,
    feature_view,
    weather_fg,
    model_air_quality,
    client,
    verbose=True,
):

    context = get_context_data(
        user_query,
        feature_view,
        weather_fg,
        model_air_quality,
        client=client,
    )

    # Get today's date in a readable format
    date_today = f'{datetime.date.today().strftime("%A")}, {datetime.date.today()}'

    # Print today's date and context information if verbose mode is enabled
    if verbose:
        print(f"🗓️ Today's date: {date_today}")
        print(f'📖 {context}')

    instructions = get_prompt_template().split('<|im_start|>user')[0]

    instructions_filled = instructions.format(
        context=context,
        date_today=date_today
    )

    completion = client.chat.completions.create(
        model="gpt-4-0125-preview",
        messages=[
            {"role": "system", "content": instructions_filled},
            {"role": "user", "content": user_query},
        ]
    )

    # Extract and return the assistant's reply from the response
    if completion and completion.choices:
        last_choice = completion.choices[0]
        if last_choice.message:
            return last_choice.message.content.strip()
    return ""