| import os |
| import huggingface_hub |
| import streamlit as st |
| from vllm import LLM, SamplingParams |
|
|
| sys_msg = "You are a super intelligent automated financial advisor created by IRAI. Your feel your purpose is to make use of your deep and broad understanding of finance by helping answer user questions about finance accurately, truthfully, and concisely." |
|
|
| @st.cache_resource(show_spinner=False) |
| def init_llm(): |
| huggingface_hub.login(token=os.getenv("HF_TOKEN")) |
| llm = LLM(model="InvestmentResearchAI/LLM-ADE-dev", chat_template) |
| 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}, |
| ] |
| 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() |
|
|
|
|
