Update legaldoc_summarizer/app.py

legaldoc_summarizer/app.py

@@ -1,3 +1,4 @@
+import os
 import streamlit as st
 from shared.hf_helpers import build_pipeline
 import yaml
@@ -10,13 +11,29 @@ def main():
     with open(CONFIG_PATH) as f:
         cfg = yaml.safe_load(f)
 
-    model_name = st.selectbox("Model:", [cfg["base_model"], "models/legaldoc_summarizer"])
+    base_model = cfg["base_model"]
+    finetuned_model = cfg.get("finetuned_model") or os.getenv("LEGALDOC_MODEL_ID")
+    local_model_dir = Path(cfg.get("finetuned_local_dir", "models/legaldoc_summarizer"))
+
+    model_options = [base_model]
+    if finetuned_model:
+        model_options.append(finetuned_model)
+    elif local_model_dir.exists():
+        model_options.append(str(local_model_dir))
+    else:
+        st.info(
+            "Using the base model until a fine-tuned checkpoint is available. "
+            "Train a model to populate `models/legaldoc_summarizer` or set `LEGALDOC_MODEL_ID` / `finetuned_model`."
+        )
+
+    model_name = st.selectbox("Model:", model_options)
+    hf_token = os.getenv("HUGGINGFACEHUB_API_TOKEN") or os.getenv("HF_TOKEN")
 
     @st.cache_resource
-    def get_pipeline(model_name):
-        return build_pipeline(model_name)
+    def get_pipeline(model_name, token):
+        return build_pipeline(model_name, token=token)
 
-    pipe = get_pipeline(model_name)
+    pipe = get_pipeline(model_name, hf_token)
 
     st.write("Paste a contract clause or judgment text below:")
     text = st.text_area("Clause or Legal Text", height=250)