generator.py

import os
import streamlit as st
from langchain_community.llms import HuggingFaceEndpoint
from pydantic import ValidationError

st.write(
	"Has environment variables been set:",
	os.environ["HUGGINGFACE_API_TOKEN"] == st.secrets["HUGGINGFACE_API_TOKEN"])


def load_llm(repo_id="mistralai/Mistral-7B-Instruct-v0.2"):
    '''
    Load the LLM from the HuggingFace model hub
    Args:
        repo_id (str): The HuggingFace model ID
    Returns:
        llm (HuggingFaceEndpoint): The LLM model
    '''

    try:
        secret_token = os.getenv("HUGGINGFACE_API_TOKEN")
        llm = HuggingFaceEndpoint(
            repo_id=repo_id,
            model_kwargs={"max_length": 128},
            api_token=secret_token
        )
        return llm
    except ValidationError as e:
        print("Validation Error:", e)
        # Log or handle the validation error appropriately
        return None
    except Exception as e:
        print("Error:", e)
        # Log or handle other exceptions
        return None