Update app.py

app.py

@@ -161,30 +161,60 @@ def extract_with_llm(text: str) -> List[str]:
         st.error(str(e))
         return []
 
+# def is_meaningful_prompt(text: str) -> bool:
+#     too_short = len(text.strip()) < 10
+#     banned_keywords = ["phone", "email", "fax", "address", "date", "contact", "website"]
+#     contains_bad_word = any(word in text.lower() for word in banned_keywords)
+#     is_just_punctuation = all(c in ":.*- " for c in text.strip())
+
+#     return not (too_short or contains_bad_word or is_just_punctuation)
 
 # =================== Format Retrieved Chunks ===================
 def format_docs(docs: List[Document]) -> str:
     return "\n\n".join(doc.page_content or doc.metadata.get("content", "") for doc in docs)
 
 # =================== Generate Response from Hugging Face Model ===================
-def generate_response(input_dict: Dict[str, Any]) -> str:
-    client = InferenceClient(api_key=HF_TOKEN.strip())
+# def generate_response(input_dict: Dict[str, Any]) -> str:
+#     client = InferenceClient(api_key=HF_TOKEN.strip())
+#     prompt = grantbuddy_prompt.format(**input_dict)
+
+#     try:
+#         response = client.chat.completions.create(
+#             model="HuggingFaceH4/zephyr-7b-beta",
+#             messages=[
+#                 {"role": "system", "content": prompt},
+#                 {"role": "user", "content": input_dict["question"]},
+#             ],
+#             max_tokens=1000,
+#             temperature=0.2,
+#         )
+#         return response.choices[0].message.content
+#     except Exception as e:
+#         st.error(f"❌ Error from model: {e}")
+#         return "⚠️ Failed to generate response. Please check your model, HF token, or request format."
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+@st.cache_resource
+def load_local_model():
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+tokenizer, model = load_local_model()
+
+def generate_response(input_dict):
     prompt = grantbuddy_prompt.format(**input_dict)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=512,
+        temperature=0.7,
+        do_sample=True
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).split("QUESTION:")[-1].strip()
 
-    try:
-        response = client.chat.completions.create(
-            model="HuggingFaceH4/zephyr-7b-beta",
-            messages=[
-                {"role": "system", "content": prompt},
-                {"role": "user", "content": input_dict["question"]},
-            ],
-            max_tokens=1000,
-            temperature=0.2,
-        )
-        return response.choices[0].message.content
-    except Exception as e:
-        st.error(f"❌ Error from model: {e}")
-        return "⚠️ Failed to generate response. Please check your model, HF token, or request format."
 
 
 # =================== RAG Chain ===================
@@ -199,13 +229,12 @@ def main():
     st.set_page_config(page_title="Grant Buddy RAG", page_icon="🤖")
     st.title("🤖 Grant Buddy: Grant-Writing Assistant")
 
+    retriever = init_vector_search().as_retriever(search_kwargs={"k": 10, "score_threshold": 0.75})
+    rag_chain = get_rag_chain(retriever)
+
     uploaded_file = st.file_uploader("Upload PDF or TXT for extra context (optional)", type=["pdf", "txt"])
     uploaded_text = ""
 
-    retriever = init_vector_search().as_retriever(search_kwargs={"k": 10, "score_threshold": 0.75})
-    rag_chain = get_rag_chain(retriever)  # ✅ Initialize before usage
-
-    # 🔁 Process uploaded file
     if uploaded_file:
         with st.spinner("📄 Processing uploaded file..."):
             if uploaded_file.name.endswith(".pdf"):
@@ -214,26 +243,40 @@ def main():
             elif uploaded_file.name.endswith(".txt"):
                 uploaded_text = uploaded_file.read().decode("utf-8")
 
-            questions = extract_with_llm(uploaded_text)
-            st.success(f"✅ Found {len(questions)} questions or headers.")
-            with st.expander("🧠 Extracted Prompts from Upload"):
-                st.write(questions)
+        # 🧠 Extract prompts using LLM
+        questions = extract_with_llm(uploaded_text)
+
+        # 🚫 Filter out irrelevant junk
+        def is_meaningful_prompt(text: str) -> bool:
+            too_short = len(text.strip()) < 10
+            banned_keywords = ["phone", "email", "fax", "address", "date", "contact", "website"]
+            contains_bad_word = any(word in text.lower() for word in banned_keywords)
+            is_just_punctuation = all(c in ":.*- " for c in text.strip())
+            return not (too_short or contains_bad_word or is_just_punctuation)
+
+        filtered_questions = [q for q in questions if is_meaningful_prompt(q)]
 
-            # Generate answers
-            answers = []
-            for q in questions:
-                full_query = f"{q}\n\nAdditional context:\n{uploaded_text}"
-                response = rag_chain.invoke(full_query)
-                answers.append({"question": q, "answer": response})
+        # 🎯 Prompt selection UI
+        selected_questions = st.multiselect("✅ Choose prompts to answer:", filtered_questions, default=filtered_questions)
+
+        if selected_questions:
+            with st.spinner("💡 Generating answers..."):
+                answers = []
+                for q in selected_questions:
+                    full_query = f"{q}\n\nAdditional context:\n{uploaded_text}"
+                    response = rag_chain.invoke(full_query)
+                    answers.append({"question": q, "answer": response})
 
             for item in answers:
                 st.markdown(f"### ❓ {item['question']}")
                 st.markdown(f"💬 {item['answer']}")
+        else:
+            st.info("No prompts selected for answering.")
 
-    # ✅ Manual query box
+    # ✍️ Manual single-question input
     query = st.text_input("Ask a grant-related question")
     if st.button("Submit"):
-        if not query and not uploaded_file:
+        if not query:
             st.warning("Please enter a question.")
             return
 
@@ -245,13 +288,14 @@ def main():
         with st.expander("🔍 Retrieved Chunks"):
             context_docs = retriever.get_relevant_documents(full_query)
             for doc in context_docs:
-                st.markdown(f"**Chunk ID:** {doc.metadata.get('chunk_id', 'unknown','title')}")
+                st.markdown(f"**Chunk ID:** {doc.metadata.get('chunk_id', 'unknown')} | **Title:** {doc.metadata.get('title', 'unknown')}")
                 st.markdown(doc.page_content[:700] + "...")
                 st.markdown("---")
 
 
 
 
+
 if __name__ == "__main__":
     main()