Update app.py

app.py

@@ -8,6 +8,7 @@
 # app.py
 import os
 import re
+from huggingface_hub import InferenceClient
 import json
 from huggingface_hub import HfApi
 import streamlit as st
@@ -130,17 +131,33 @@ def init_vector_search() -> MongoDBAtlasVectorSearch:
 #     return headers + questions
 def extract_with_llm(text: str) -> List[str]:
     client = InferenceClient(api_key=HF_TOKEN.strip())
-    prompt = (
-        "Extract a list of grant application headers and questions from the following text. "
-        "Include section titles, prompts, or any questions that ask for a response. "
-        "Return them as a numbered list.\n\n"
-        f"{text[:3000]}"  # Optional: limit input to avoid token overflow
-    )
     try:
-        response = client.text_generation(prompt=prompt, max_new_tokens=500)
-        return [line.strip("-•1234567890. ") for line in response.split("\n") if line.strip()]
+        response = client.chat.completions.create(
+            model="mistralai/Mistral-Nemo-Instruct-2407",  # or "HuggingFaceH4/zephyr-7b-beta"
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are an assistant helping extract questions and headers from grant applications.",
+                },
+                {
+                    "role": "user",
+                    "content": (
+                        "Please extract all the grant application headers and questions from the following text. "
+                        "Include section titles, prompts, and any question-like content. Return them as a numbered list.\n\n"
+                        f"{text[:3000]}"
+                    ),
+                },
+            ],
+            temperature=0.2,
+            max_tokens=512,
+        )
+        return [
+            line.strip("•-1234567890. ").strip()
+            for line in response.choices[0].message.content.strip().split("\n")
+            if line.strip()
+        ]
     except Exception as e:
-        st.error("❌ Failed to extract questions with LLM")
+        st.error("❌ LLM extraction failed")
         st.error(str(e))
         return []