import os
import huggingface_hub
import streamlit as st
from vllm import LLM, SamplingParams

sys_msg = """You are an expert financial advisor named IRAI. You have a comprehensive understanding of finance and investing with experience and expertise in all areas of finance. 
#Objective:
Answer questions accurately and truthfully given your current knowledge.  You do not have access to up-to-date current market data; this will be available in the future. Answer the question directly.
Style and tone:
Answer in a friendly and engaging manner representing a top female investment professional working at a leading investment bank.
#Audience:
The questions will be asked by top technology executives and CFO of large fintech companies and successful startups.
#Response:
Direct answer to question, concise yet insightful."""

@st.cache_resource(show_spinner="Loading model..")
def init_llm():
    huggingface_hub.login(token=os.getenv("HF_TOKEN"))
    llm = LLM(model="InvestmentResearchAI/LLM-ADE-dev")
    tok = llm.get_tokenizer()
    tok.eos_token = '<|im_end|>' # Override to use turns
    return llm

def get_response(prompt):
    try:
        convo = [
            {"role": "system", "content": sys_msg},
            {"role": "user", "content": prompt},
        ]
        llm = init_llm()
        prompts = [llm.get_tokenizer().apply_chat_template(convo, tokenize=False)]
        sampling_params = SamplingParams(temperature=0.3, top_p=0.95, max_tokens=500, stop_token_ids=[128009])
        outputs = llm.generate(prompts, sampling_params)
        for output in outputs:
            return output.outputs[0].text
    except Exception as e:
        return f"An error occurred: {str(e)}"


def main():
    st.title("LLM-ADE 9B Demo")
    
    input_text = st.text_area("Enter your text here:", value="", height=200)
    
    if st.button("Generate"):
        if input_text:
            with st.spinner('Generating response...'):
                response_text = get_response(input_text)
                st.write(response_text)
        else:
            st.warning("Please enter some text to generate a response.")

llm = init_llm()

if __name__ == "__main__":
    main()