Update app.py

app.py

@@ -2,24 +2,20 @@ import streamlit as st
 import pandas as pd
 from langchain_core.documents import Document
 from langchain.prompts import PromptTemplate
-from langchain_community.llms import LlamaCpp  # To run the GGUF model
-from langchain_community.vectorstores import FAISS
-from langchain.text_splitter import MarkdownHeaderTextSplitter
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_huggingface import HuggingFaceEndpoint
 import json
-import hashlib
 
 # --- Streamlit Page Configuration ---
 st.set_page_config(
-    page_title="URS document extractor",
+    page_title="URS Document Extractor",
     page_icon="📄",
     layout="wide"
 )
 
-st.title("📄 URS Document Data Extractor")
-st.markdown("This app extracts key information from a hardcoded URS document context using Llama 3 on Hugging Face Spaces.")
+st.title("📄 URS Document Data Extractor (API Version)")
+st.markdown("This app extracts key information from a URS document using the fast and reliable Hugging Face Inference API.")
 
-# --- Core Logic from your script ---
+# --- Core Logic ---
 
 # Using the hardcoded context you provided
 CONTEXT_DOCUMENTS = [
@@ -43,25 +39,21 @@ FINAL_QUERY = {
     "[Requirement number_MOC]": "Which section number corresponds to the slope value of the desired level of instruments and accuracy levels?"
 }
 
-# Use Streamlit's cache to load the model only once
+# --- Function to initialize the LLM ---
+# This uses the Hugging Face Inference API. It's fast and reliable.
 @st.cache_resource
 def get_llm():
-    """Loads the Llama 3 model from Hugging Face."""
-    st.info("Downloading and loading the Llama 3 model... This may take a few minutes on first run.")
-    
-    # This function will download the model from the Hugging Face Hub and cache it.
-    llm = LlamaCpp(
-        # The repo ID of the model you want to use
-        model_path_or_repo_id="QuantFactory/Meta-Llama-3-8B-Instruct-GGUF",
-        # The specific model file in that repository
-        model_basename="Meta-Llama-3-8B-Instruct.Q4_K_M.gguf",
-        n_gpu_layers=-1, # Offload all layers to GPU if available (won't use on basic HF space)
-        n_batch=512,
-        n_ctx=4096,      # Context window
-        f16_kv=True,     # Must be set to True on metal for Apple silicon
-        verbose=True,
-        temperature=0,
-        max_tokens=2048
+    """Initializes the HuggingFaceEndpoint for the Llama 3 model."""
+    # Ensure the secret is available before initializing
+    if "HF_TOKEN" not in st.secrets:
+        st.error("Hugging Face token not found! Please add it to your Space secrets.", icon="🚨")
+        st.stop()
+        
+    llm = HuggingFaceEndpoint(
+        repo_id="meta-llama/Meta-Llama-3-8B-Instruct",
+        huggingfacehub_api_token=st.secrets["HF_TOKEN"],
+        temperature=0.1,
+        max_new_tokens=2048,
     )
     return llm
 
@@ -103,39 +95,41 @@ def generating_results(context, Final_Query, llm):
     return response
 
 # --- Streamlit App UI ---
+if 'show_button' not in st.session_state:
+    st.session_state.show_button = True
+
+# We need the secret to run, so we check for it first.
+if "HF_TOKEN" not in st.secrets:
+    st.error("A Hugging Face Token is required to run this app. Please add it to your Space secrets in the 'Settings' tab.", icon="🔑")
+    st.session_state.show_button = False
 
-st.info("Click the button below to start the extraction process. Note: This uses a large model on a free CPU, so it will be slow and may take several minutes.")
-
-if st.button("✨ Start Extraction", type="primary"):
-    with st.spinner("Please wait... The Llama 3 model is processing the document..."):
-        try:
-            # 1. Load the LLM
-            # This is cached, so it's fast after the first run.
-            llm_extraction = get_llm()
-            
-            # 2. Call the generation function
-            # The context is already hardcoded in this example.
-            # In a real app, you would generate it from a file upload.
-            raw_response = generating_results(CONTEXT_DOCUMENTS, FINAL_QUERY, llm_extraction)
-            
-            st.subheader("Raw Model Output")
-            st.text(raw_response)
-
-            # 3. Safely parse the JSON from the response
+if st.session_state.show_button:
+    if st.button("✨ Start Extraction", type="primary"):
+        with st.spinner("Calling the Llama 3 Inference API... This should be fast!"):
             try:
-                # The model is instructed to return a JSON block. Find it.
-                json_start = raw_response.find('{')
-                json_end = raw_response.rfind('}') + 1
-                json_str = raw_response[json_start:json_end]
-                result = json.loads(json_str)
+                # 1. Load the LLM via the API
+                llm_extraction = get_llm()
                 
-                st.success("✅ Extraction Complete!")
-                st.subheader("Parsed JSON Result")
-                st.json(result)
+                # 2. Call the generation function
+                raw_response = generating_results(CONTEXT_DOCUMENTS, FINAL_QUERY, llm_extraction)
+                
+                st.subheader("Raw Model Output")
+                st.text(raw_response)
+
+                # 3. Safely parse the JSON from the response
+                try:
+                    json_start = raw_response.find('{')
+                    json_end = raw_response.rfind('}') + 1
+                    json_str = raw_response[json_start:json_end]
+                    result = json.loads(json_str)
+                    
+                    st.success("✅ Extraction Complete!")
+                    st.subheader("Parsed JSON Result")
+                    st.json(result)
+
+                except Exception as e:
+                    st.error(f"⚠️ Failed to parse JSON from the model's response. Error: {e}")
+                    st.text_area("Model's Raw Response for Debugging:", raw_response, height=300)
 
             except Exception as e:
-                st.error(f"⚠️ Failed to parse JSON from the model's response. Error: {e}")
-                st.text_area("Model's Raw Response for Debugging:", raw_response, height=300)
-
-        except Exception as e:
-            st.error(f"An unexpected error occurred during the extraction process: {e}")
+                st.error(f"An unexpected error occurred during the extraction process: {e}")