Update financegpt/app.py

financegpt/app.py

@@ -1,28 +1,31 @@
 import streamlit as st
 from shared.hf_helpers import build_pipeline
 import yaml
+from pathlib import Path
 
-st.set_page_config(page_title="FinanceGPT Demo", page_icon="💰", layout="centered")
-
-st.title("💰 FinanceGPT — Financial Q&A Demo")
+def main():
+    st.set_page_config(page_title="FinanceGPT Demo", page_icon="💰", layout="centered")
+    st.title("💰 FinanceGPT — Financial Q&A Demo")
 
-from pathlib import Path
-CONFIG_PATH = Path(__file__).resolve().parent / "config.yaml"
-with open(CONFIG_PATH) as f:
-    cfg = yaml.safe_load(f)
+    CONFIG_PATH = Path(__file__).resolve().parent / "config.yaml"
+    with open(CONFIG_PATH) as f:
+        cfg = yaml.safe_load(f)
 
+    model_name = st.selectbox("Select model:", [cfg["base_model"], "models/financegpt"])
 
-model_name = st.selectbox("Select model:", [cfg["base_model"], "models/financegpt"])
+    @st.cache_resource
+    def get_pipe(model_name):
+        return build_pipeline(model_name)
 
-@st.cache_resource
-def get_pipe(model_name):
-    return build_pipeline(model_name)
+    pipe = get_pipe(model_name)
 
-pipe = get_pipe(model_name)
+    prompt = st.text_area("Enter a financial statement or question:")
+    if st.button("Generate Answer"):
+        if prompt.strip():
+            result = pipe(prompt, max_new_tokens=cfg["demo"]["max_new_tokens"])
+            st.markdown("### 🧠 Answer")
+            st.write(result[0]["generated_text"])
 
-prompt = st.text_area("Enter a financial statement or question:")
-if st.button("Generate Answer"):
-    if prompt.strip():
-        result = pipe(prompt, max_new_tokens=cfg["demo"]["max_new_tokens"])
-        st.markdown("### 🧠 Answer")
-        st.write(result[0]["generated_text"])
+# 🔁 So it works standalone or via streamlit_hub.py
+if __name__ == "__main__":
+    main()