Update app.py

app.py

@@ -1,216 +1,115 @@
-import os
-import io
-import re
-import json
-import base64
-import requests
-import numpy as np
+import os, io, re, json, base64, requests, numpy as np
 import streamlit as st
 from pypdf import PdfReader
 import matplotlib.pyplot as plt
 
 # -----------------------------
-# Config / Secrets (safe)
+# Config
 # -----------------------------
 st.set_page_config(page_title="PDF Summarizer + Audio + QA", page_icon="📄", layout="wide")
 
-# Prefer environment variable (Spaces sets secrets as env vars), *then* try st.secrets safely.
 HF_TOKEN = os.environ.get("HF_TOKEN", "")
-if not HF_TOKEN:
-    try:
-        # Access st.secrets inside try/except so we don't crash when no secrets file exists.
-        HF_TOKEN = st.secrets.get("HF_TOKEN", "") if hasattr(st, "secrets") else ""
-    except Exception:
-        HF_TOKEN = ""
-
 HEADERS_JSON = {
     "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
     "Content-Type": "application/json",
     "Accept": "application/json",
 }
 
-SUMMARIZER_MODEL = "facebook/bart-large-cnn"
-TTS_MODEL = "facebook/mms-tts-eng"
+SUMMARIZER_MODEL = "pszemraj/long-t5-tglobal-base-16384-book-summary"
+TTS_MODELS = [
+    "espnet/kan-bayashi_ljspeech_vits",
+    "facebook/fastspeech2-en-ljspeech"
+]
 EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 QA_MODEL = "deepset/roberta-base-squad2"
 
 # -----------------------------
-# Helper: Hugging Face inference
+# API helpers
 # -----------------------------
-def hf_infer_json(model_id: str, payload: dict, router=False, accept=None, timeout=120):
-    """
-    Send request to Hugging Face Hosted Inference API.
-    If `router=True` we'll use the router base path (useful for some pipelines).
-    If backend returns binary (audio), this returns raw bytes.
-    """
+def hf_infer_json(model_id: str, payload: dict, router=False, accept=None):
     if router:
         url = f"https://router.huggingface.co/hf-inference/models/{model_id}"
     else:
         url = f"https://api-inference.huggingface.co/models/{model_id}"
-
     headers = HEADERS_JSON.copy()
     if accept:
         headers["Accept"] = accept
-
-    try:
-        r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
-        r.raise_for_status()
-    except requests.exceptions.RequestException as e:
-        # Bubble up a useful message
-        raise RuntimeError(f"Hugging Face request failed: {e}")
-
-    # Try to decode JSON; if fails, return bytes/content
+    r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=120)
+    r.raise_for_status()
     try:
         return r.json()
-    except ValueError:
+    except requests.exceptions.JSONDecodeError:
         return r.content
 
-# -----------------------------
-# Text / PDF utilities
-# -----------------------------
-def extract_text_from_pdf(file) -> str:
-    reader = PdfReader(file)
-    pages = []
-    for p in reader.pages:
-        try:
-            pages.append(p.extract_text() or "")
-        except Exception:
-            pages.append("")
-    return "\n".join(pages)
-
-def clean_text(s: str) -> str:
-    return re.sub(r"\s+", " ", s).strip()
-
-def split_into_chunks(text: str, max_chars: int = 1800, overlap: int = 200):
-    text = clean_text(text)
+def split_into_chunks(text: str, max_chars: int = 1500, overlap: int = 200):
+    text = re.sub(r"\s+", " ", text).strip()
     chunks = []
     i = 0
     while i < len(text):
         chunk = text[i:i+max_chars]
         last_dot = chunk.rfind(". ")
         if last_dot > 400:
-            chunk = chunk[: last_dot + 1]
+            chunk = chunk[:last_dot+1]
             i += last_dot + 1 - overlap
         else:
             i += max_chars - overlap
         chunks.append(chunk.strip())
     return [c for c in chunks if c]
 
-# -----------------------------
-# Embeddings + similarity
-# -----------------------------
 def embed_texts(texts):
-    """
-    Calls the feature-extraction pipeline on the router endpoint.
-    Returns numpy array shape (n_texts, dim)
-    """
     url = f"https://router.huggingface.co/hf-inference/models/{EMB_MODEL}/pipeline/feature-extraction"
-    headers = {
-        "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
-        "Content-Type": "application/json",
-        "Accept": "application/json",
-    }
-    try:
-        r = requests.post(url, headers=headers, data=json.dumps({"inputs": texts}), timeout=120)
-        r.raise_for_status()
-    except requests.exceptions.RequestException as e:
-        raise RuntimeError(f"Embedding request failed: {e}")
-
+    headers = HEADERS_JSON
+    r = requests.post(url, headers=headers, data=json.dumps({"inputs": texts}), timeout=120)
+    r.raise_for_status()
     arr = np.array(r.json(), dtype=np.float32)
-
-    # Cases:
-    # - arr.ndim == 1 -> single vector (dim,) -> reshape to (1,dim)
-    # - arr.ndim == 2 -> batch of vectors (n, dim) -> return as-is
-    # - arr.ndim == 3 -> model returned token-level vectors per item: mean-pool per item -> (n, dim)
-    if arr.ndim == 1:
-        return arr.reshape(1, -1)
     if arr.ndim == 2:
-        return arr
+        return arr.mean(axis=0, keepdims=True)
     if arr.ndim == 3:
-        pooled = np.array([a.mean(axis=0) for a in arr])
-        return pooled
-    # Fallback
-    return arr.reshape(arr.shape[0], -1)
+        pooled = [a.mean(axis=0) for a in arr]
+        return np.vstack(pooled)
+    return np.array(arr)
 
 def cosine_sim(a, b):
-    """
-    a: (m, d), b: (n, d) -> returns (m, n)
-    """
-    a_n = a / (np.linalg.norm(a, axis=-1, keepdims=True) + 1e-8)
-    b_n = b / (np.linalg.norm(b, axis=-1, keepdims=True) + 1e-8)
-    return a_n @ b_n.T
+    a = a / (np.linalg.norm(a, axis=-1, keepdims=True) + 1e-8)
+    b = b / (np.linalg.norm(b, axis=-1, keepdims=True) + 1e-8)
+    return a @ b.T
 
-# -----------------------------
-# Summarization
-# -----------------------------
-def summarize_long_text(text: str, per_chunk_max_len=220, final_max_len=250):
-    chunks = split_into_chunks(text, max_chars=1800, overlap=200)
+def summarize_long_text(text: str):
+    chunks = split_into_chunks(text)
     mini_summaries = []
     for c in chunks:
-        try:
-            out = hf_infer_json(
-                SUMMARIZER_MODEL,
-                {"inputs": c, "parameters": {"max_length": per_chunk_max_len, "min_length": 60, "do_sample": False}},
-                router=False,
-            )
-        except Exception as e:
-            # if API fails, include the chunk (truncated) as fallback
-            mini_summaries.append(c[:1000])
-            continue
-
-        # Hosted inference often returns a list of dicts with 'summary_text'
-        if isinstance(out, list) and len(out) and isinstance(out[0], dict) and "summary_text" in out[0]:
+        out = hf_infer_json(SUMMARIZER_MODEL, {"inputs": c}, router=False)
+        if isinstance(out, list) and len(out) and "summary_text" in out[0]:
             mini_summaries.append(out[0]["summary_text"])
-        elif isinstance(out, dict) and "summary_text" in out:
-            mini_summaries.append(out["summary_text"])
         else:
-            mini_summaries.append(c[:1000])
-
-    joined = " ".join(mini_summaries)
-    try:
-        final = hf_infer_json(
-            SUMMARIZER_MODEL,
-            {"inputs": joined, "parameters": {"max_length": final_max_len, "min_length": 80, "do_sample": False}},
-            router=False,
-        )
-    except Exception:
-        return joined[:1200], chunks
-
-    if isinstance(final, list) and len(final) and isinstance(final[0], dict) and "summary_text" in final[0]:
-        return final[0]["summary_text"], chunks
-    if isinstance(final, dict) and "summary_text" in final:
-        return final["summary_text"], chunks
+            mini_summaries.append(c[:800])
+    return " ".join(mini_summaries), chunks
 
-    return joined[:1200], chunks
-
-# -----------------------------
-# TTS
-# -----------------------------
 def tts_wav_bytes(text: str) -> bytes:
-    try:
-        res = hf_infer_json(TTS_MODEL, {"inputs": text}, router=False, accept="audio/wav", timeout=180)
-    except Exception as e:
-        raise RuntimeError(f"TTS request failed: {e}")
-
-    if isinstance(res, (bytes, bytearray)):
-        return res
-    if isinstance(res, dict) and "audio" in res:
+    for model in TTS_MODELS:
         try:
-            return base64.b64decode(res["audio"])
+            res = hf_infer_json(model, {"inputs": text}, router=False, accept="audio/wav")
+            if isinstance(res, (bytes, bytearray)):
+                return res
+            if isinstance(res, dict) and "audio" in res:
+                return base64.b64decode(res["audio"])
         except Exception:
-            pass
-    raise RuntimeError("TTS API did not return audio bytes.")
+            continue
+    raise RuntimeError("All TTS models failed.")
+
+def extract_text_from_pdf(file) -> str:
+    reader = PdfReader(file)
+    pages = []
+    for p in reader.pages:
+        try:
+            pages.append(p.extract_text() or "")
+        except:
+            pages.append("")
+    return "\n".join(pages)
 
-# -----------------------------
-# Visualization helper
-# -----------------------------
 def make_word_freq_chart(text: str, top_k=20):
     text = text.lower()
-    stop = set(
-        (
-            "the a an and of to in is are for with on by as at this that from be was were it its it's into or if not your you we they their our can may such more most other also than which".split()
-        )
-    )
+    stop = set(("the a an and of to in is are for with on by as at this that from be was were it its it’s into or if not your you we they their our can may such more most other also than which".split()))
     tokens = re.findall(r"[a-zA-Z]{3,}", text)
     freq = {}
     for t in tokens:
@@ -233,14 +132,10 @@ def make_word_freq_chart(text: str, top_k=20):
 # UI
 # -----------------------------
 st.title("📄 PDF → Summary · 🔊 Audio · 📊 Chart · ❓ Q&A")
-st.caption("Powered by Hugging Face Hosted Inference API (free models).")
-
-if not HF_TOKEN:
-    st.warning("No HF_TOKEN found. Add HF_TOKEN in Space Settings → Secrets (recommended). The app will still run but HF API calls will fail without a token.")
+st.caption("Free models via Hugging Face Hosted Inference API.")
 
 uploaded = st.file_uploader("Upload a PDF", type=["pdf"])
 
-# session state
 if "doc_text" not in st.session_state:
     st.session_state.doc_text = ""
     st.session_state.chunks = []
@@ -254,22 +149,17 @@ if uploaded:
         st.success(f"Loaded {len(text)} characters.")
 
     st.write("### Actions")
-    c1, c2, c3 = st.columns(3)
-
-    with c1:
+    with st.container():
         if st.button("📝 Summarize"):
             with st.spinner("Summarizing..."):
-                try:
-                    summary, chunks = summarize_long_text(st.session_state.doc_text)
-                    st.session_state.summary = summary
-                    st.session_state.chunks = chunks
-                    st.success("Summary ready.")
-                    st.write("#### Summary")
-                    st.write(summary)
-                except Exception as e:
-                    st.error(f"Summarization failed: {e}")
-
-    with c2:
+                summary, chunks = summarize_long_text(st.session_state.doc_text)
+                st.session_state.summary = summary
+                st.session_state.chunks = chunks
+            st.success("Summary ready.")
+            st.write("#### Summary")
+            st.write(st.session_state.summary)
+
+    with st.container():
         if st.button("🔊 Generate Audio (summary)"):
             target_text = st.session_state.summary or st.session_state.doc_text[:1200]
             with st.spinner("Generating audio..."):
@@ -280,7 +170,7 @@ if uploaded:
                 except Exception as e:
                     st.error(f"TTS failed: {e}")
 
-    with c3:
+    with st.container():
         if st.button("📊 Show Word-Frequency Chart"):
             with st.spinner("Building chart..."):
                 make_word_freq_chart(st.session_state.doc_text)
@@ -293,10 +183,11 @@ if uploaded:
             st.session_state.chunks = split_into_chunks(st.session_state.doc_text)
         with st.spinner("Thinking..."):
             try:
-                # embed once/cache
                 if st.session_state.chunk_vecs is None:
-                    st.session_state.chunk_vecs = embed_texts(st.session_state.chunks)
-                vecs = st.session_state.chunk_vecs
+                    vecs = embed_texts(st.session_state.chunks)
+                    st.session_state.chunk_vecs = vecs
+                else:
+                    vecs = st.session_state.chunk_vecs
 
                 q_vec = embed_texts([question])
                 sims = cosine_sim(q_vec, vecs).flatten()