app.py

import json 
import os

import requests
import streamlit as st

CHATBOT_ENDPOINT = os.environ["CHATBOT_ENDPOINT"]
TOKEN = os.environ["TOKEN"]

def generate(prompt):

    try:
        r = requests.post(
            CHATBOT_ENDPOINT,
            data=json.dumps({"instruction": prompt, "token": TOKEN}),
            headers = {"content-type": "application/json"},
            timeout = 20,
            stream = True,
        )
        new_str = b""
        for s in r.iter_content():
            new_str += s
            try:
                output_str = new_str.decode("utf-8")
                yield output_str.replace("\n", "  \n")
                new_str = b""
            except:
                pass
    except:
        yield "<<Some errors occured>>"

intro = """This is a demo site for Stockmark-LLM-100b. This service is running on AWS Inferentia2.
- Pretrained model: [stockmark/stockmark-100b](https://huggingface.co/stockmark/stockmark-100b)
- Instruction tuned model: [stockmark/stockmark-100b-instruct-v0.1](https://huggingface.co/stockmark/stockmark-100b-instruct-v0.1)
"""

disclaimer = """
- Responses of our LLM may be incorrect, biased, or harmful.
- We may use users chat data in this demo to improve our LLM.
"""


tab1, tab2, tab3 = st.tabs(["Demo", "Disclaimer"])




with tab1:
    st.title("Stockmark-LLM-100b")
    
    
    
    st.markdown(intro)
    
    prompt = st.session_state.get("prompt", "")
    response = st.session_state.get("response", "")
    
    if prompt == "" or response:
        print("new_session")
        prompt_new = st.text_area("Prompt:")
        if prompt_new:
            st.session_state["prompt"] = prompt_new
            st.session_state["response"] = ""
            st.rerun()
    else:
        prompt = st.text_area("Prompt:", value=prompt, disabled=True)
    
    if prompt:
    
        if response:
            with st.chat_message("assistant"):
                st.write(response)
        else:
            with st.chat_message("assistant"):
                response = st.write_stream(generate(prompt))
            
            st.session_state["response"] = response
            st.rerun()

with tab2:
    st.title("Stockmark-LLM-100b: Disclaimer")
    st.markdown(disclaimer)