|
import os |
|
import huggingface_hub |
|
import streamlit as st |
|
from vllm import LLM, SamplingParams |
|
|
|
sys_msg = """#Context: |
|
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. |
|
Style and tone: |
|
Please 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: |
|
Answer, 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|>' |
|
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() |
|
|
|
|