Create app.py

app.py

@@ -0,0 +1,163 @@
+import streamlit as st
+from io import BytesIO
+import ibm_watsonx_ai
+import secretsload
+import anton_ego_jimmy
+import requests
+import time
+import re
+
+from ibm_watsonx_ai.foundation_models import ModelInference
+from ibm_watsonx_ai import Credentials, APIClient
+from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
+from ibm_watsonx_ai.metanames import GenTextReturnOptMetaNames as RetParams
+from secretsload import load_stsecrets # script to load credentials from HuggingFace secrets section
+
+credentials = load_stsecrets()
+
+st.set_page_config(
+    page_title="Jimmy le critique",
+    page_icon="🍷",
+    initial_sidebar_state="collapsed"
+)
+
+# Password protection
+def check_password():
+    def password_entered():
+        if st.session_state["password"] == st.secrets["app_password"]:
+            st.session_state["password_correct"] = True
+            del st.session_state["password"]
+        else:
+            st.session_state["password_correct"] = False
+
+    if "password_correct" not in st.session_state:
+        st.markdown("\n\n")
+        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
+        st.divider()
+        st.info("Template of Jimmy developed by Milan Mrdenovic © IBM Norway 2024")
+        return False
+    elif not st.session_state["password_correct"]:
+        st.markdown("\n\n")
+        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
+        st.divider()
+        st.info("Template of Jimmy developed by Milan Mrdenovic © IBM Norway 2024")
+        st.error("😕 Password incorrect")
+        return False
+    else:
+        return True
+
+if not check_password():
+    st.stop()
+
+
+# Initialize session state
+if 'current_page' not in st.session_state:
+    st.session_state.current_page = 0
+
+def initialize_session_state():
+    if 'chat_history' not in st.session_state:
+        st.session_state.chat_history = []
+
+def setup_client():
+    credentials = Credentials(
+        url=st.secrets["url"],
+        api_key=st.secrets["api_key"]
+    )
+    apo = st.secrets["api_key"]
+    client = APIClient(credentials, project_id=st.secrets["project_id"])
+    return client
+
+client = setup_client()
+
+def prepare_prompt(prompt, chat_history):
+    if anton_ego_jimmy.TYPE == "chat" and chat_history:
+        chats = "\n".join([f"{message['role']}: \"{message['content']}\"" for message in chat_history])
+        return f"Conversation History:\n{chats}\n\nNew Message: {prompt}"
+    return prompt
+
+def apply_prompt_syntax(prompt, system_prompt, prompt_template, bake_in_prompt_syntax):
+    model_family_syntax = {
+        "llama3-instruct (llama-3 & 3.1) - system": """\n<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n""",
+        "llama3-instruct (llama-3 & 3.1) - user": """\n<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n""",
+        "granite-13b-chat & instruct - system": """\n<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>\n\n""",
+        "granite-13b-chat & instruct - user": """\n<|user|>\n{prompt}\n<|assistant|>\n\n""",
+        "llama2-chat - system": """\n[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n{prompt} [/INST]\n""",
+        "llama2-chat - user": """\n[INST] {prompt} [/INST] """,
+        "mistral & mixtral v2 tokenizer - system": """\n<s>[INST]System Prompt:[{system_prompt}]\n\n{prompt} [/INST]\n""",
+        "mistral & mixtral v2 tokenizer - system segmented": """\n<s>[INST]System Prompt:{system_prompt}[/INST][INST]{prompt} [/INST]\n""",
+        "mistral & mixtral v2 tokenizer - user": """\n<s>[INST]{prompt} [/INST]\n"""
+    }
+    
+    if bake_in_prompt_syntax:
+        template = model_family_syntax[prompt_template]
+        if system_prompt:
+            return template.format(system_prompt=system_prompt, prompt=prompt)
+    return prompt
+
+def generate_response(watsonx_llm, prompt_data, params):
+    generated_response = watsonx_llm.generate_text_stream(prompt=prompt_data, params=params)
+    for chunk in generated_response:
+        yield chunk
+
+def get_response(user_input):
+    # Prepare the prompt
+    prompt = prepare_prompt(user_input, st.session_state.chat_history)
+
+    # Apply prompt syntax
+    prompt_data = apply_prompt_syntax(
+        prompt, 
+        anton_ego_jimmy.SYSTEM_PROMPT,
+        anton_ego_jimmy.PROMPT_TEMPLATE,
+        anton_ego_jimmy.BAKE_IN_PROMPT_SYNTAX
+    )
+
+    watsonx_llm = ModelInference(
+        api_client=client, 
+        model_id=anton_ego_jimmy.SELECTED_MODEL,
+        verify=anton_ego_jimmy.VERIFY
+    )
+
+    # Prepare parameters
+    params = {
+        GenParams.DECODING_METHOD: anton_ego_jimmy.DECODING_METHOD,
+        GenParams.MAX_NEW_TOKENS: anton_ego_jimmy.MAX_NEW_TOKENS,
+        GenParams.MIN_NEW_TOKENS: anton_ego_jimmy.MIN_NEW_TOKENS,
+        # GenParams.LENGTH_PENALTY: anton_ego_jimmy.LENGTH_PENALTY,
+        GenParams.REPETITION_PENALTY: anton_ego_jimmy.REPETITION_PENALTY,
+        GenParams.STOP_SEQUENCES: anton_ego_jimmy.STOP_SEQUENCES
+    }
+
+    print(prompt_data, params)
+
+    # Generate and stream response
+    with st.chat_message("assistant", avatar="🧐"):
+        stream = generate_response(watsonx_llm, prompt_data, params)
+        response = st.write_stream(stream)
+
+    # Add AI response to chat history
+    st.session_state.chat_history.append({"role": "assistant", "content": response})
+
+    return response
+
+def main():
+    initialize_session_state()
+    st.title("Jimmy le critique 🧐")
+
+    # User input
+    if anton_ego_jimmy.DISPLAY_CHAT_HISTORY == 1:
+        for message in st.session_state.chat_history:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+    user_input = st.chat_input("You:", key="user_input")
+
+    if user_input:
+        # Add user message to chat history
+        st.session_state.chat_history.append({"role": "user", "content": user_input})
+        with st.chat_message("user", avatar="🥷🏻"):
+            st.markdown(user_input)
+
+        # Get response
+        get_response(user_input)
+
+if __name__ == "__main__":
+    main()