Refactored HF token ID access to a common function

pages/005_LLM_Models.py

@@ -11,7 +11,6 @@ if st_setup('LLM Models'):
         HFLlamaChatModel.load_configs()
 
     SESSION_KEY_CHAT_SERVER = 'chat_server'
-    HF_AUTH_KEY_SECRET = 'hf_token'
     button_count = 0
 
 
@@ -56,7 +55,7 @@ if st_setup('LLM Models'):
                 with chat_container:
                     with st.chat_message("user"):
                         st.write(prompt)
-                    chat_model = HFLlamaChatModel.get_model(st.session_state[SESSION_KEY_CHAT_SERVER])
-                    response = chat_model(prompt, st.secrets[HF_AUTH_KEY_SECRET])
+                    chat_model = HFLlamaChatModel.for_name(st.session_state[SESSION_KEY_CHAT_SERVER])
+                    response = chat_model(prompt)
                     with st.chat_message("assistant"):
                         st.write(response)

src/common.py

@@ -1,5 +1,13 @@
 import os
+import streamlit as st
 
 
 data_dir = os.path.join(os.path.dirname(__file__), '..', 'data')
 config_dir = os.path.join(os.path.dirname(__file__), '..', 'config')
+
+def hf_api_token() -> str:
+    #TODO: Need to consider how to make this more generic to look for a token elsewhere
+    token = st.secrets['hf_token']
+    if token is None:
+        raise ValueError('No HF access token found in streamlit secrets')
+    return token

src/models.py

@@ -3,7 +3,7 @@ import os
 import requests
 from typing import List
 
-from src.common import config_dir
+from src.common import config_dir, hf_api_token
 
 
 class HFLlamaChatModel:
@@ -16,13 +16,13 @@ class HFLlamaChatModel:
             configs = json.load(f)['models']
             cls.models = []
             for cfg in configs:
-                if cls.get_model(cfg['name']) is None:
+                if cls.for_name(cfg['name']) is None:
                     cls.models.append(HFLlamaChatModel(cfg['name'], cfg['id'], cfg['description']))
 
     @classmethod
-    def get_model(cls, model: str):
+    def for_name(cls, name: str):
         for m in cls.models:
-            if m.name == model:
+            if m.name == name:
                 return m
 
     @classmethod
@@ -38,10 +38,12 @@ class HFLlamaChatModel:
 
     def __call__(self,
                  query: str,
-                 auth_token: str,
+                 auth_token: str = None,
                  system_prompt: str = None,
                  max_new_tokens: str = 256,
                  temperature: float = 1.0):
+        if auth_token is None:
+            auth_token = hf_api_token()  # Attempt look up if not provided
         headers = {"Authorization": f"Bearer {auth_token}"}
         api_url = f"https://api-inference.huggingface.co/models/{self.id}"
         if system_prompt is None: