added BAM support

app.py

@@ -17,12 +17,13 @@ from langchain import OpenAI
 from langchain.document_loaders import WebBaseLoader, TextLoader, Docx2txtLoader, PyMuPDFLoader
 from whatsapp_chat_custom import WhatsAppChatLoader # use this instead of from langchain.document_loaders import WhatsAppChatLoader
 
-from ibm_watson_machine_learning.foundation_models.utils.enums import ModelTypes
 from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams
 from ibm_watson_machine_learning.foundation_models.utils.enums import DecodingMethods
 from ibm_watson_machine_learning.foundation_models import Model
 from ibm_watson_machine_learning.foundation_models.extensions.langchain import WatsonxLLM
 
+import genai
+
 from collections import deque
 import re
 from bs4 import BeautifulSoup
@@ -64,28 +65,37 @@ if mode.type!='userInputDocs':
 
 ###############################################################################################
 
-def setOaiApiKey(api_key):
-    credComps = [oaiKey_btn, wxKey_tb, wxPid_tb, wxKey_btn]
-    api_key = getOaiCreds(api_key)
+def setOaiApiKey(creds):
+    creds = getOaiCreds(creds)
     try:
-        openai.Model.list(api_key=api_key.get('oai_key','Null')) # test the API key
-        api_key_st = api_key
-        return oaiKey_tb.update('API Key accepted', interactive=False, type='text'), *[x.update(interactive=False) for x in credComps], api_key_st
+        openai.Model.list(api_key=creds.get('oai_key','Null')) # test the API key
+        api_key_st = creds
+        return 'OpenAI credentials accepted', *[x.update(interactive=False) for x in credComps_btn_tb], api_key_st
     except Exception as e:
-        return oaiKey_tb.update(str(e), type='text'), *[x.update() for x in credComps+[api_key_state]]
-    
+        gr.Warning(str(e))
+        return [x.update() for x in credComps_op]
+
+def setBamApiKey(creds):
+    creds = getBamCreds(creds)
+    try:
+        genai.Model.models(credentials=creds['bam_creds'])
+        api_key_st = creds
+        return 'BAM credentials accepted', *[x.update(interactive=False) for x in credComps_btn_tb], api_key_st
+    except Exception as e:
+        gr.Warning(str(e))
+        return [x.update() for x in credComps_op]
 
 def setWxApiKey(key, p_id):
-    credComps = [wxKey_btn, oaiKey_tb, oaiKey_btn]
-    api_key = getWxCreds(key, p_id)
+    creds = getWxCreds(key, p_id)
     try:
-        testModel = Model(model_id=ModelTypes.FLAN_UL2, credentials=api_key['credentials'], project_id=api_key['project_id']) # test the API key
-        del testModel
-        api_key_st = api_key
-        return *[x.update('Watsonx credentials accepted', interactive=False, type='text') for x in [wxKey_tb, wxPid_tb]], *[x.update(interactive=False) for x in credComps], api_key_st
+        Model(model_id='google/flan-ul2', credentials=creds['credentials'], project_id=creds['project_id']) # test the API key
+        api_key_st = creds
+        return 'Watsonx credentials accepted', *[x.update(interactive=False) for x in credComps_btn_tb], api_key_st
     except Exception as e:
-        return *[x.update(str(e), type='text') for x in [wxKey_tb, wxPid_tb]], *[x.update() for x in credComps+[api_key_state]]
+        gr.Warning(str(e))
+        return [x.update() for x in credComps_op]
     
+
 # convert user uploaded data to vectorstore
 def uiData_vecStore(userFiles, userUrls, api_key_st, vsDict_st={}, progress=gr.Progress()):
     opComponents = [data_ingest_btn, upload_fb, urls_tb]
@@ -102,6 +112,7 @@ def uiData_vecStore(userFiles, userUrls, api_key_st, vsDict_st={}, progress=gr.P
         for file in file_paths:
             os.remove(file)
     else:
+        gr.Error('No documents found')
         return {}, '', *[x.update() for x in opComponents]
     # Splitting and Chunks
     docs = split_docs(documents)
@@ -109,7 +120,8 @@ def uiData_vecStore(userFiles, userUrls, api_key_st, vsDict_st={}, progress=gr.P
     try:
         embeddings = getEmbeddingFunc(api_key_st)
     except Exception as e:
-        return {}, str(e), *[x.update() for x in opComponents]
+        gr.Error(str(e))
+        return {}, '', *[x.update() for x in opComponents]
     
     progress(0.5, 'Creating Vector Database')
     vsDict_st = getVsDict(embeddings, docs, vsDict_st)
@@ -130,45 +142,30 @@ def initializeChatbot(temp, k, modelName, stdlQs, api_key_st, vsDict_st, progres
     if mode.welcomeMsg:
         welMsg = mode.welcomeMsg
     else:
-        welMsg = qa_chain_st({'question': initialize_prompt, 'chat_history':[]})['answer']
+        # welMsg = qa_chain_st({'question': initialize_prompt, 'chat_history':[]})['answer']
+        welMsg = welcomeMsgDefault
     print('Chatbot initialized at ', datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
 
     return qa_chain_st, chainTuple[1], btn.update(interactive=True), initChatbot_btn.update('Chatbot ready. Now visit the chatbot Tab.', interactive=False)\
-        , oaiKey_tb.update(), gr.Tabs.update(selected='cb'), chatbot.update(value=[('', welMsg)])
+        , status_tb.update(), gr.Tabs.update(selected='cb'), chatbot.update(value=[('Hi', welMsg)])
 
 # just update the QA Chain, no updates to any UI
 def updateQaChain(temp, k, modelNameDD, stdlQs, api_key_st, vsDict_st):
     # if we are not adding data from ui, then use vsDict_hard as vectorstore
     if vsDict_st=={} and mode.type!='userInputDocs': vsDict_st=vsDict_hard
     
-    if api_key_st.get('service')=='openai':
+    if api_key_st['service']=='openai':
         if not 'openai' in modelNameDD:
             modelNameDD = 'gpt-3.5-turbo (openai)'  # default model for openai
-        modelName = modelNameDD.split('(')[0].strip()
-        # check if the input model is chat model or legacy model
-        try:
-            ChatOpenAI(openai_api_key=api_key_st.get('oai_key','Null'), temperature=0,model_name=modelName,max_tokens=1).predict('')
-            llm = ChatOpenAI(openai_api_key=api_key_st.get('oai_key','Null'), temperature=float(temp),model_name=modelName)
-        except:
-            OpenAI(openai_api_key=api_key_st.get('oai_key','Null'), temperature=0,model_name=modelName,max_tokens=1).predict('')
-            llm = OpenAI(openai_api_key=api_key_st.get('oai_key','Null'), temperature=float(temp),model_name=modelName)
-    elif api_key_st.get('service')=='watsonx':
+        llm = getOaiLlm(temp, modelNameDD, api_key_st)
+    elif api_key_st['service']=='watsonx':
         if not 'watsonx' in modelNameDD:
             modelNameDD = 'meta-llama/llama-2-70b-chat (watsonx)' # default model for watsonx
-        modelName = modelNameDD.split('(')[0].strip()
-        wxModelParams = {
-            GenParams.DECODING_METHOD: DecodingMethods.SAMPLE,
-            GenParams.MAX_NEW_TOKENS: 1000,
-            GenParams.MIN_NEW_TOKENS: 1,
-            GenParams.TEMPERATURE: float(temp),
-            GenParams.TOP_K: 50,
-            GenParams.TOP_P: 1
-        }
-        flan_ul2_model = Model(
-                model_id=modelName, 
-                params=wxModelParams, 
-                credentials=api_key_st['credentials'], project_id=api_key_st['project_id'])
-        llm = WatsonxLLM(model=flan_ul2_model)
+        llm = getWxLlm(temp, modelNameDD, api_key_st)
+    elif api_key_st['service']=='bam':
+        if not 'bam' in modelNameDD:
+            modelNameDD = 'ibm/granite-13b-sft (bam)' # default model for bam
+        llm = getBamLlm(temp, modelNameDD, api_key_st)
     else:
         raise Exception('Error: Invalid or None Credentials')
     # settingsUpdated = 'Settings updated:'+ ' Model=' + modelName + ', Temp=' + str(temp)+ ', k=' + str(k)
@@ -196,7 +193,7 @@ def updateQaChain(temp, k, modelNameDD, stdlQs, api_key_st, vsDict_st):
         
 
 def respond(message, chat_history, qa_chain):
-    result = qa_chain({'question': message, "chat_history": [tuple(x) for x in chat_history]})
+    result = qa_chain({'question': message, "chat_history": [tuple(x) for x in chat_history[1:]]})
     src_docs = getSourcesFromMetadata([x.metadata for x in result["source_documents"]], sourceOnly=False)[0]
     # streaming
     streaming_answer = ""
@@ -227,6 +224,10 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue='orange', secondary_hue='gray
                     oaiKey_tb = gr.Textbox(label="OpenAI API Key", type='password'\
                             , info='You can find OpenAI API key at https://platform.openai.com/account/api-keys')
                     oaiKey_btn = gr.Button("Submit OpenAI API Key")
+                with gr.Column():
+                    bamKey_tb = gr.Textbox(label="BAM API Key", type='password'\
+                            , info='Internal IBMers only')
+                    bamKey_btn = gr.Button("Submit BAM API Key")
                 with gr.Column():
                     wxKey_tb = gr.Textbox(label="Watsonx API Key", type='password'\
                             , info='You can find IBM Cloud API Key at Manage > Access (IAM) > API keys on https://cloud.ibm.com/iam/overview')
@@ -239,12 +240,15 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue='orange', secondary_hue='gray
                                     , info=url_tb_info\
                                     , placeholder=url_tb_ph)
                 data_ingest_btn = gr.Button("Load Data")
-            status_tb = gr.TextArea(label='Status bar', show_label=False, visible=mode.uiAddDataVis)
+            status_tb = gr.TextArea(label='Status Info')
             initChatbot_btn = gr.Button("Initialize Chatbot", variant="primary")
 
+        credComps_btn_tb = [oaiKey_tb, oaiKey_btn, bamKey_tb, bamKey_btn, wxKey_tb, wxPid_tb, wxKey_btn]
+        credComps_op = [status_tb] + credComps_btn_tb + [api_key_state]
+        
         with gr.Tab('Chatbot', id='cb'):
             with gr.Row():
-                chatbot = gr.Chatbot(label="Chat History", scale=2)
+                chatbot = gr.Chatbot(label="Chat History", scale=2, avatar_images=(user_avatar, bot_avatar))
                 srcDocs = gr.TextArea(label="References")
             msg = gr.Textbox(label="User Input",placeholder="Type your questions here")
             with gr.Row():
@@ -266,12 +270,17 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue='orange', secondary_hue='gray
     ### Setup the Gradio Event Listeners
 
     # OpenAI API button
-    oaiKey_btn_args = {'fn':setOaiApiKey, 'inputs':[oaiKey_tb], 'outputs':[oaiKey_tb, oaiKey_btn, wxKey_tb, wxPid_tb, wxKey_btn, api_key_state]}
+    oaiKey_btn_args = {'fn':setOaiApiKey, 'inputs':[oaiKey_tb], 'outputs':credComps_op}
     oaiKey_btn.click(**oaiKey_btn_args)
     oaiKey_tb.submit(**oaiKey_btn_args)
 
+    # BAM API button
+    bamKey_btn_args = {'fn':setBamApiKey, 'inputs':[bamKey_tb], 'outputs':credComps_op}
+    bamKey_btn.click(**bamKey_btn_args)
+    bamKey_tb.submit(**bamKey_btn_args)
+
     # Watsonx Creds button
-    wxKey_btn_args = {'fn':setWxApiKey, 'inputs':[wxKey_tb, wxPid_tb], 'outputs':[wxKey_tb, wxPid_tb, wxKey_btn, oaiKey_tb, oaiKey_btn, api_key_state]}
+    wxKey_btn_args = {'fn':setWxApiKey, 'inputs':[wxKey_tb, wxPid_tb], 'outputs':credComps_op}
     wxKey_btn.click(**wxKey_btn_args)
 
     # Data Ingest Button

requirements.txt

@@ -9,4 +9,5 @@ PyMuPDF
 gdown
 docx2txt
 sentence-transformers
-ibm-watson-machine-learning
+ibm-watson-machine-learning
+ibm-generative-ai

ttyd_consts.py

@@ -8,6 +8,11 @@ initialize_prompt = """Write a short welcome message to the user. Describe the d
  If this data is about a person, mention his name instead of using pronouns. After describing the overview, you should mention top 3 example questions that the user can ask about this data.\
  \n\nYour response should be short and precise. Format of your response should be Summary:\n{Description and Summary} \n\n Example Questions:\n{Example Questions}"""
 
+
+user_avatar = 'https://cdn-icons-png.flaticon.com/512/6861/6861326.png'
+# user_avatar = None
+bot_avatar = 'https://cdn-icons-png.flaticon.com/512/1782/1782384.png'
+
 nustian_exps = ['Tell me about NUSTIAN',
                 'Who is the NUSTIAN regional lead for Silicon Valley?',
                 'Tell me details about NUSTIAN coaching program.',
@@ -23,10 +28,35 @@ stdlQs_rb_choices =  ['Retrieve relavant docs using original question, send orig
                     , 'Retrieve relavant docs using standalone question, send standalone question to LLM']
 
 
-
-model_dd_info = 'You can also input any OpenAI model name, compatible with /v1/completions or /v1/chat/completions endpoint. Details: https://platform.openai.com/docs/models/'
-
-model_dd_choices = ['gpt-3.5-turbo (openai)', 'gpt-3.5-turbo-16k (openai)', 'gpt-4 (openai)', 'text-davinci-003 (Legacy - openai)', 'text-curie-001 (Legacy - openai)', 'babbage-002 (openai)'] + [model.value+' (watsonx)' for model in ModelTypes]
+bam_models = sorted(['bigscience/bloom',
+ 'salesforce/codegen2-16b',
+ 'codellama/codellama-34b-instruct',
+ 'tiiuae/falcon-40b',
+ 'ibm/falcon-40b-8lang-instruct',
+ 'google/flan-t5-xl',
+ 'google/flan-t5-xxl',
+ 'google/flan-ul2',
+ 'eleutherai/gpt-neox-20b',
+ 'togethercomputer/gpt-neoxt-chat-base-20b',
+ 'ibm/granite-13b-sft',
+ 'ibm/granite-13b-sft-cft',
+ 'ibm/granite-3b-code-v1',
+ 'meta-llama/llama-2-13b',
+ 'meta-llama/llama-2-13b-chat',
+ 'meta-llama/llama-2-13b-chat-beam',
+ 'meta-llama/llama-2-70b',
+ 'meta-llama/llama-2-70b-chat',
+ 'meta-llama/llama-2-7b',
+ 'meta-llama/llama-2-7b-chat',
+ 'mosaicml/mpt-30b',
+ 'ibm/mpt-7b-instruct',
+ 'bigscience/mt0-xxl',
+ 'bigcode/starcoder',
+ 'google/ul2'])
+
+model_dd_info = 'You can also input any OpenAI model name or BAM model ID.'
+
+model_dd_choices = ['gpt-3.5-turbo (openai)', 'gpt-3.5-turbo-16k (openai)', 'gpt-4 (openai)', 'text-davinci-003 (Legacy - openai)', 'text-curie-001 (Legacy - openai)', 'babbage-002 (openai)'] + [model.value+' (watsonx)' for model in ModelTypes] + [model + ' (bam)' for model in bam_models]
 
 url_tb_info = 'Upto 100 domain webpages will be crawled for each URL. You can also enter online PDF files.'
 
@@ -70,6 +100,7 @@ welcomeMsgArslan = """Summary: The document provides a comprehensive overview of
 3. Tell me about Arslan's educational background.
 """
 
+welcomeMsgDefault = """Hello and welcome! I'm your personal data assistant. Ask me anything about your data and I'll try my best to answer."""
 
 class TtydMode():
     def __init__(self, name='', title='', type='', dir=None, files=[], urls=[], vis=False, welMsg='', def_k=4):

ttyd_functions.py

@@ -20,6 +20,19 @@ import mimetypes
 from pathlib import Path
 import tiktoken
 
+from langchain.chat_models import ChatOpenAI
+from langchain import OpenAI
+
+from ibm_watson_machine_learning.metanames import GenTextParamsMetaNames as GenParams
+from ibm_watson_machine_learning.foundation_models.utils.enums import DecodingMethods
+from ibm_watson_machine_learning.foundation_models import Model
+from ibm_watson_machine_learning.foundation_models.extensions.langchain import WatsonxLLM
+
+
+import genai
+from genai.extensions.langchain import LangChainInterface
+from genai.schemas import GenerateParams
+
 # Regex pattern to match a URL
 HTTP_URL_PATTERN = r'^http[s]*://.+'
 
@@ -28,21 +41,26 @@ media_files = tuple([x for x in mimetypes.types_map if mimetypes.types_map[x].sp
 filter_strings = ['/email-protection#']
 
 def getOaiCreds(key):
-    if key:
-        return {'service': 'openai',
-                    'oai_key' : key
-                }
-    else:
-        return {} 
+    key = key if key else 'Null'
+    return {'service': 'openai',
+                'oai_key' : key
+            }
+
+
+def getBamCreds(key):
+    key = key if key else 'Null'
+    return {'service': 'bam',
+                'bam_creds' : genai.Credentials(key, api_endpoint='https://bam-api.res.ibm.com/v1')
+            }
+
 
 def getWxCreds(key, p_id):
-    if key and p_id:
-        return {'service': 'watsonx',
+    key = key if key else 'Null'
+    p_id = p_id if p_id else 'Null'
+    return {'service': 'watsonx',
                 'credentials' : {"url": "https://us-south.ml.cloud.ibm.com", "apikey": key },
-                'project_id': p_id
-                }
-    else:
-        return {}
+                    'project_id': p_id
+            }
 
 def getPersonalBotApiKey():
     if os.getenv("OPENAI_API_KEY"):
@@ -52,6 +70,45 @@ def getPersonalBotApiKey():
     else:
         return {}
     
+
+
+def getOaiLlm(temp, modelNameDD, api_key_st):
+    modelName = modelNameDD.split('(')[0].strip()
+    # check if the input model is chat model or legacy model
+    try:
+        ChatOpenAI(openai_api_key=api_key_st['oai_key'], temperature=0,model_name=modelName,max_tokens=1).predict('')
+        llm = ChatOpenAI(openai_api_key=api_key_st['oai_key'], temperature=float(temp),model_name=modelName)
+    except:
+        OpenAI(openai_api_key=api_key_st['oai_key'], temperature=0,model_name=modelName,max_tokens=1).predict('')
+        llm = OpenAI(openai_api_key=api_key_st['oai_key'], temperature=float(temp),model_name=modelName)
+    return llm
+
+
+def getWxLlm(temp, modelNameDD, api_key_st):
+    modelName = modelNameDD.split('(')[0].strip()
+    wxModelParams = {
+        GenParams.DECODING_METHOD: DecodingMethods.SAMPLE,
+        GenParams.MAX_NEW_TOKENS: 1000,
+        GenParams.MIN_NEW_TOKENS: 1,
+        GenParams.TEMPERATURE: float(temp),
+        GenParams.TOP_K: 50,
+        GenParams.TOP_P: 1
+    }
+    model = Model(
+            model_id=modelName, 
+            params=wxModelParams, 
+            credentials=api_key_st['credentials'], project_id=api_key_st['project_id'])
+    llm = WatsonxLLM(model=model)
+    return llm
+
+
+def getBamLlm(temp, modelNameDD, api_key_st):
+    modelName = modelNameDD.split('(')[0].strip()
+    parameters = GenerateParams(decoding_method="sample", max_new_tokens=1024, temperature=float(temp), top_k=50, top_p=1)
+    llm = LangChainInterface(model=modelName, params=parameters, credentials=api_key_st['bam_creds'])
+    return llm
+
+
 def get_hyperlinks(url):
     try:
         reqs = requests.get(url)
@@ -249,7 +306,7 @@ def getEmbeddingFunc(creds):
         if creds.get('service')=='openai':
             embeddings = OpenAIEmbeddings(openai_api_key=creds.get('oai_key','Null'))
         # WX key used
-        elif creds.get('service')=='watsonx':
+        elif creds.get('service')=='watsonx' or creds.get('service')=='bam':
             # testModel = Model(model_id=ModelTypes.FLAN_UL2, credentials=creds['credentials'], project_id=creds['project_id']) # test the API key
             # del testModel
             embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") # for now use OpenSource model for embedding as WX doesnt have any embedding model