app.py

import os
import huggingface_hub
import streamlit as st
from config import config
from vllm import LLM, SamplingParams
from functioncall import ModelInference

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"), new_session=False)
    llm = ModelInference(chat_template='chatml')
    return llm

def get_response(prompt):
    try:
        return llm.generate_function_call(
            prompt, 
            config.chat_template, 
            config.num_fewshot, 
            config.max_depth
        )
    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()

def main_headless():
    while True:
        input_text = input("Enter your text here: ")
        print(get_response(input_text))

if __name__ == "__main__":
    main_headless()