Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,6 +7,7 @@ from typing import List, Dict, Any
 import numpy as np
 import streamlit as st
 import PyPDF2
+import requests
 from dotenv import load_dotenv
 from huggingface_hub import InferenceClient, login
 from streamlit_chat import message as st_message
@@ -77,6 +78,7 @@ def find_handbook() -> List[str]:
     st.error("❌ No PDF found in the same folder as this app.")
     return []
 
+
 def load_pdf_texts(pdf_paths: List[str]) -> List[Dict[str, Any]]:
     """Extract text from all pages of provided PDFs."""
     pages = []
@@ -89,6 +91,7 @@ def load_pdf_texts(pdf_paths: List[str]) -> List[Dict[str, Any]]:
                     pages.append({"filename": os.path.basename(path), "page": i + 1, "text": text})
     return pages
 
+
 def chunk_text(pages: List[Dict[str, Any]], size: int, overlap: int) -> List[Dict[str, Any]]:
     """Split text into overlapping chunks."""
     chunks = []
@@ -106,20 +109,45 @@ def chunk_text(pages: List[Dict[str, Any]], size: int, overlap: int) -> List[Dic
             start += size - overlap
     return chunks
 
+
 def embed_texts(texts: List[str]) -> np.ndarray:
-    """Get embeddings via Hugging Face Inference API."""
-    if not hf_client:
-        st.error("❌ No Hugging Face client initialized.")
+    """Get embeddings via Hugging Face Inference API with fallback."""
+    if not HF_TOKEN:
+        st.error("❌ Missing HF_TOKEN.")
         return np.zeros((len(texts), 768))
+
+    # --- Primary method: InferenceClient.feature_extraction ---
     try:
-        emb = hf_client.post(
-            "/embeddings",
-            json={"inputs": texts, "model": EMBED_MODEL},
+        embeddings = hf_client.feature_extraction(
+            model=EMBED_MODEL,
+            inputs=texts
         )
-        return np.array(emb["embeddings"])
-    except Exception as e:
-        st.error(f"Embedding error: {e}")
-        return np.zeros((len(texts), 768))
+
+        # Handle nested list outputs (token-level vectors)
+        if isinstance(embeddings[0][0], list):
+            embeddings = [np.mean(np.array(e), axis=0) for e in embeddings]
+
+        return np.array(embeddings)
+
+    # --- Fallback method: REST API ---
+    except Exception as e1:
+        st.warning(f"⚠️ feature_extraction() failed, using REST API fallback: {e1}")
+        try:
+            headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+            response = requests.post(
+                f"https://api-inference.huggingface.co/pipeline/feature-extraction/{EMBED_MODEL}",
+                headers=headers,
+                json={"inputs": texts}
+            )
+            response.raise_for_status()
+            data = response.json()
+
+            embeddings = [np.mean(np.array(e), axis=0) for e in data]
+            return np.array(embeddings)
+        except Exception as e2:
+            st.error(f"Embedding error: {e2}")
+            return np.zeros((len(texts), 768))
+
 
 def build_faiss_index(chunks: List[Dict[str, Any]]) -> None:
     """Build and save FAISS index for handbook chunks."""
@@ -139,6 +167,7 @@ def build_faiss_index(chunks: List[Dict[str, Any]]) -> None:
     with open(EMB_DIM_FILE, "w") as f:
         json.dump({"dim": dim}, f)
 
+
 def load_faiss_index():
     """Load FAISS index and metadata if available."""
     if not (os.path.exists(INDEX_FILE) and os.path.exists(META_FILE)):
@@ -148,6 +177,7 @@ def load_faiss_index():
         meta = json.load(f)
     return index, meta
 
+
 def search_index(query: str, index, meta, top_k: int, threshold: float) -> List[Dict[str, Any]]:
     """Search FAISS for top-K similar chunks."""
     query_emb = embed_texts([query])
@@ -160,6 +190,7 @@ def search_index(query: str, index, meta, top_k: int, threshold: float) -> List[
             results.append(result)
     return results
 
+
 def generate_answer(context: str, query: str) -> str:
     """Generate robust answer with explicit citations — auto-switches between endpoints."""
     prompt = f"""
@@ -210,6 +241,7 @@ If the answer is not explicitly found, respond with:
         except Exception as e2:
             return f"⚠️ Error generating answer: {e2}"
 
+
 # =============================================================
 # 🔍 Index Handling
 # =============================================================
@@ -233,6 +265,7 @@ def ensure_index():
         st.stop()
     return index, meta
 
+
 # =============================================================
 # 💬 Chat Interface
 # =============================================================