Allow users to provide their own HF access token/API key

app.py

@@ -54,19 +54,6 @@ def _get_prompt_template(is_refinement: bool) -> str:
     return template
 
 
-@st.cache_resource
-def _get_llm(repo_id: str, max_new_tokens: int):
-    """
-    Get an LLM instance.
-
-    :param repo_id: The model name.
-    :param max_new_tokens: The max new tokens to generate.
-    :return: The LLM.
-    """
-
-    return llm_helper.get_hf_endpoint(repo_id, max_new_tokens)
-
-
 APP_TEXT = _load_strings()
 
 # Session variables
@@ -81,18 +68,35 @@ texts = list(GlobalConfig.PPTX_TEMPLATE_FILES.keys())
 captions = [GlobalConfig.PPTX_TEMPLATE_FILES[x]['caption'] for x in texts]
 
 with st.sidebar:
+    # The PPT templates
     pptx_template = st.sidebar.radio(
-        'Select a presentation template:',
+        '1: Select a presentation template:',
         texts,
         captions=captions,
         horizontal=True
     )
-    st.divider()
-    llm_to_use = st.sidebar.selectbox(
-        'Select an LLM to use:',
-        [f'{k} ({v["description"]})' for k, v in GlobalConfig.HF_MODELS.items()]
+
+    # The LLMs
+    llm_provider_to_use = st.sidebar.selectbox(
+        label='2: Select an LLM to use:',
+        options=[f'{k} ({v["description"]})' for k, v in GlobalConfig.VALID_MODELS.items()],
+        index=0,
+        help=(
+            'LLM provider codes:\n\n'
+            '- **[hf]**: Hugging Face Inference Endpoint\n'
+        ),
     ).split(' ')[0]
 
+    # The API key/access token
+    api_key_token = st.text_input(
+        label=(
+            '3: Paste your API key/access token:\n\n'
+            '*Optional* if an HF Mistral LLM is selected from the list but still encouraged.\n\n'
+        ),
+        type='password',
+    )
+    st.caption('(Wrong HF access token will lead to validation error)')
+
 
 def build_ui():
     """
@@ -101,9 +105,9 @@ def build_ui():
 
     st.title(APP_TEXT['app_name'])
     st.subheader(APP_TEXT['caption'])
-    st.markdown(
-        '![Visitors](https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fbarunsaha%2Fslide-deck-ai&countColor=%23263759)'  # noqa: E501
-    )
+    # st.markdown(
+    #     '![Visitors](https://api.visitorbadge.io/api/visitors?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fbarunsaha%2Fslide-deck-ai&countColor=%23263759)'  # noqa: E501
+    # )
 
     with st.expander('Usage Policies and Limitations'):
         st.text(APP_TEXT['tos'] + '\n\n' + APP_TEXT['tos2'])
@@ -162,9 +166,15 @@ def set_up_chat_ui():
             )
             return
 
+        provider, llm_name = llm_helper.get_provider_model(llm_provider_to_use)
+
+        if not provider or not llm_name:
+            st.error('No valid LLM provider and/or model name found!')
+            return
+
         logger.info(
             'User input: %s | #characters: %d | LLM: %s',
-            prompt, len(prompt), llm_to_use
+            prompt, len(prompt), llm_name
         )
         st.chat_message('user').write(prompt)
 
@@ -193,15 +203,17 @@ def set_up_chat_ui():
         response = ''
 
         try:
-            for chunk in _get_llm(
-                    repo_id=llm_to_use,
-                    max_new_tokens=GlobalConfig.HF_MODELS[llm_to_use]['max_new_tokens']
+            for chunk in llm_helper.get_langchain_llm(
+                    provider=provider,
+                    model=llm_name,
+                    max_new_tokens=GlobalConfig.VALID_MODELS[llm_provider_to_use]['max_new_tokens'],
+                    api_key=api_key_token.strip(),
             ).stream(formatted_template):
                 response += chunk
 
                 # Update the progress bar
                 progress_percentage = min(
-                    len(response) / GlobalConfig.HF_MODELS[llm_to_use]['max_new_tokens'], 0.95
+                    len(response) / GlobalConfig.VALID_MODELS[llm_provider_to_use]['max_new_tokens'], 0.95
                 )
                 progress_bar.progress(
                     progress_percentage,

global_config.py

@@ -17,12 +17,13 @@ class GlobalConfig:
     A data class holding the configurations.
     """
 
-    HF_MODELS = {
-        'mistralai/Mistral-7B-Instruct-v0.2': {
+    VALID_PROVIDERS = {'hf'}
+    VALID_MODELS = {
+        '[hf]mistralai/Mistral-7B-Instruct-v0.2': {
             'description': 'faster, shorter',
             'max_new_tokens': 8192
         },
-        'mistralai/Mistral-Nemo-Instruct-2407': {
+        '[hf]mistralai/Mistral-Nemo-Instruct-2407': {
             'description': 'longer response',
             'max_new_tokens': 12228
         },

helpers/llm_helper.py

@@ -1,4 +1,7 @@
 import logging
+import re
+from typing import Tuple, Union
+
 import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util import Retry
@@ -9,7 +12,8 @@ from langchain_core.language_models import LLM
 from global_config import GlobalConfig
 
 
-HF_API_HEADERS = {"Authorization": f"Bearer {GlobalConfig.HUGGINGFACEHUB_API_TOKEN}"}
+LLM_PROVIDER_MODEL_REGEX = re.compile(r'\[(.*?)\](.*)')
+HF_API_HEADERS = {'Authorization': f'Bearer {GlobalConfig.HUGGINGFACEHUB_API_TOKEN}'}
 REQUEST_TIMEOUT = 35
 
 logger = logging.getLogger(__name__)
@@ -27,12 +31,31 @@ http_session.mount('https://', adapter)
 http_session.mount('http://', adapter)
 
 
-def get_hf_endpoint(repo_id: str, max_new_tokens: int) -> LLM:
+def get_provider_model(provider_model: str) -> Tuple[str, str]:
+    """
+    Parse and get LLM provider and model name from strings like `[provider]model/name-version`.
+
+    :param provider_model: The provider, model name string from `GlobalConfig`.
+    :return: The provider and the model name.
+    """
+
+    match = LLM_PROVIDER_MODEL_REGEX.match(provider_model)
+
+    if match:
+        inside_brackets = match.group(1)
+        outside_brackets = match.group(2)
+        return inside_brackets, outside_brackets
+
+    return '', ''
+
+
+def get_hf_endpoint(repo_id: str, max_new_tokens: int, api_key: str = '') -> LLM:
     """
     Get an LLM via the HuggingFaceEndpoint of LangChain.
 
     :param repo_id: The model name.
     :param max_new_tokens: The max new tokens to generate.
+    :param api_key: [Optional] Hugging Face access token.
     :return: The HF LLM inference endpoint.
     """
 
@@ -46,82 +69,54 @@ def get_hf_endpoint(repo_id: str, max_new_tokens: int) -> LLM:
         temperature=GlobalConfig.LLM_MODEL_TEMPERATURE,
         repetition_penalty=1.03,
         streaming=True,
-        huggingfacehub_api_token=GlobalConfig.HUGGINGFACEHUB_API_TOKEN,
+        huggingfacehub_api_token=api_key or GlobalConfig.HUGGINGFACEHUB_API_TOKEN,
         return_full_text=False,
         stop_sequences=['</s>'],
     )
 
 
-# def hf_api_query(payload: dict) -> dict:
-#     """
-#     Invoke HF inference end-point API.
-#
-#     :param payload: The prompt for the LLM and related parameters.
-#     :return: The output from the LLM.
-#     """
-#
-#     try:
-#         response = http_session.post(
-#             HF_API_URL,
-#             headers=HF_API_HEADERS,
-#             json=payload,
-#             timeout=REQUEST_TIMEOUT
-#         )
-#         result = response.json()
-#     except requests.exceptions.Timeout as te:
-#         logger.error('*** Error: hf_api_query timeout! %s', str(te))
-#         result = []
-#
-#     return result
-
-
-# def generate_slides_content(topic: str) -> str:
-#     """
-#     Generate the outline/contents of slides for a presentation on a given topic.
-#
-#     :param topic: Topic on which slides are to be generated.
-#     :return: The content in JSON format.
-#     """
-#
-#     with open(GlobalConfig.SLIDES_TEMPLATE_FILE, 'r', encoding='utf-8') as in_file:
-#         template_txt = in_file.read().strip()
-#         template_txt = template_txt.replace('<REPLACE_PLACEHOLDER>', topic)
-#
-#     output = hf_api_query({
-#         'inputs': template_txt,
-#         'parameters': {
-#             'temperature': GlobalConfig.LLM_MODEL_TEMPERATURE,
-#             'min_length': GlobalConfig.LLM_MODEL_MIN_OUTPUT_LENGTH,
-#             'max_length': GlobalConfig.LLM_MODEL_MAX_OUTPUT_LENGTH,
-#             'max_new_tokens': GlobalConfig.LLM_MODEL_MAX_OUTPUT_LENGTH,
-#             'num_return_sequences': 1,
-#             'return_full_text': False,
-#             # "repetition_penalty": 0.0001
-#         },
-#         'options': {
-#             'wait_for_model': True,
-#             'use_cache': True
-#         }
-#     })
-#
-#     output = output[0]['generated_text'].strip()
-#     # output = output[len(template_txt):]
-#
-#     json_end_idx = output.rfind('```')
-#     if json_end_idx != -1:
-#         # logging.debug(f'{json_end_idx=}')
-#         output = output[:json_end_idx]
-#
-#     logger.debug('generate_slides_content: output: %s', output)
-#
-#     return output
+def get_langchain_llm(
+        provider: str,
+        model: str,
+        max_new_tokens: int,
+        api_key: str = ''
+) -> Union[LLM, None]:
+    """
+    Get an LLM based on the provider and model specified.
+
+    :param provider: The LLM provider. Valid values are `hf` for Hugging Face.
+    :param model:
+    :param max_new_tokens:
+    :param api_key:
+    :return:
+    """
+    if not provider or not model or provider not in GlobalConfig.VALID_PROVIDERS:
+        return None
+
+    if provider == 'hf':
+        logger.debug('Getting LLM via HF endpoint: %s', model)
+
+        return HuggingFaceEndpoint(
+            repo_id=model,
+            max_new_tokens=max_new_tokens,
+            top_k=40,
+            top_p=0.95,
+            temperature=GlobalConfig.LLM_MODEL_TEMPERATURE,
+            repetition_penalty=1.03,
+            streaming=True,
+            huggingfacehub_api_token=api_key or GlobalConfig.HUGGINGFACEHUB_API_TOKEN,
+            return_full_text=False,
+            stop_sequences=['</s>'],
+        )
+
+    return None
 
 
 if __name__ == '__main__':
-    # results = get_related_websites('5G AI WiFi 6')
-    #
-    # for a_result in results.results:
-    #     print(a_result.title, a_result.url, a_result.extract)
+    inputs = [
+        '[hf]mistralai/Mistral-7B-Instruct-v0.2',
+        '[gg]gemini-1.5-flash-002'
+    ]
 
-    # get_ai_image('A talk on AI, covering pros and cons')
-    pass
+    for text in inputs:
+        print(get_provider_model(text))