Update app.py

app.py

@@ -1,28 +1,35 @@
-import streamlit as st
+# app.py
+
+from fastapi import FastAPI, UploadFile, File
+from pydantic import BaseModel
+from typing import List
+import fitz  # PyMuPDF
 from transformers import pipeline
-from langchain_community.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
+from sentence_transformers import SentenceTransformer
+from langchain.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_core.documents import Document as LCDocument
-import PyPDF2
-from docx import Document as DocxDocument
-import io
-from typing import List
-from pydantic import BaseModel
-import tempfile
+from langchain.schema import Document
+from langchain.chains.question_answering import load_qa_chain
+from langchain.llms import HuggingFacePipeline
+from langchain_core.documents import Document as LangchainDocument
 
+# --- Init FastAPI ---
+app = FastAPI()
 
-# === Summarizer ===
+# --- Summarizer ---
 summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
-# === QA Model ===
-qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
+# --- Question Answering ---
+qa_pipe = pipeline("question-answering", model="deepset/roberta-base-squad2")
 
-# === Embedding model ===
-embedding_function = HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
+# --- Embedding model ---
+embedding_model = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-small-en-v1.5")
 
+# --- Text Splitter ---
+text_splitter = RecursiveCharacterTextSplitter(chunk_size=800, chunk_overlap=100)
 
-# === Pydantic Models ===
+# --- Pydantic schemas ---
 class Summary(BaseModel):
     summary: str
 
@@ -33,78 +40,62 @@ class DocumentAnalysis(BaseModel):
     summary: Summary
     key_points: List[KeyPoint]
 
+class QARequest(BaseModel):
+    question: str
+    context: str
+
+class QAResponse(BaseModel):
+    answer: str
 
-# === Loaders ===
-def load_pdf(file):
-    reader = PyPDF2.PdfReader(file)
+# --- PDF Text Extractor ---
+def extract_text_from_pdf(pdf_file: UploadFile) -> str:
     text = ""
-    for page in reader.pages:
-        text += page.extract_text()
+    with fitz.open(stream=pdf_file.file.read(), filetype="pdf") as doc:
+        for page in doc:
+            text += page.get_text()
     return text
 
-def load_docx(file):
-    doc = DocxDocument(file)
-    return "\n".join([para.text for para in doc.paragraphs])
-
-
-# === Analysis ===
-def analyze_text_structured(text):
-    result = summarizer(text, max_length=200, min_length=50, do_sample=False)[0]["summary_text"]
-    key_points = [KeyPoint(point=line.strip()) for line in result.split(". ") if line.strip()]
-    return DocumentAnalysis(summary=Summary(summary=result), key_points=key_points)
-
-# === Embedding & Retrieval ===
-def get_vectorstore_from_text(text):
-    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
-    chunks = splitter.split_text(text)
-    docs = [LCDocument(page_content=chunk) for chunk in chunks]
-    return FAISS.from_documents(docs, embedding_function)
-
-def answer_question(vectorstore, question):
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-    docs = retriever.get_relevant_documents(question)
-    context = "\n".join([doc.page_content for doc in docs])
-    result = qa_pipeline(question=question, context=context)
+# --- Analyze Text (summarization) ---
+def analyze_text_structured(text: str) -> DocumentAnalysis:
+    chunks = text_splitter.split_text(text)
+    summaries = []
+    for chunk in chunks:
+        result = summarizer(chunk, max_length=200, min_length=50, do_sample=False)
+        if result:
+            summaries.append(result[0]["summary_text"])
+
+    full_summary = " ".join(summaries)
+    key_points = [KeyPoint(point=line.strip()) for line in full_summary.split(". ") if line.strip()]
+    return DocumentAnalysis(summary=Summary(summary=full_summary), key_points=key_points)
+
+# --- Question Answering ---
+def answer_question(question: str, context: str) -> str:
+    result = qa_pipe(question=question, context=context)
     return result["answer"]
 
-
-# === Streamlit UI ===
-st.title("📄 AI Document Analyzer")
-
-uploaded_file = st.file_uploader("Upload a document (PDF or DOCX)", type=["pdf", "docx"])
-input_text = st.text_area("Or paste your text here", height=200)
-
-if st.button("Analyze"):
-    if uploaded_file:
-        file_bytes = uploaded_file.read()
-        file_ext = uploaded_file.name.split(".")[-1]
-        if file_ext == "pdf":
-            text = load_pdf(io.BytesIO(file_bytes))
-        elif file_ext == "docx":
-            text = load_docx(io.BytesIO(file_bytes))
-        else:
-            st.error("Unsupported file format.")
-            st.stop()
-    elif input_text:
-        text = input_text
-    else:
-        st.warning("Please upload a file or paste text.")
-        st.stop()
-
-    with st.spinner("Analyzing..."):
-        analysis = analyze_text_structured(text)
-        vectorstore = get_vectorstore_from_text(text)
-
-    st.subheader("🔍 Summary")
-    st.write(analysis.summary.summary)
-
-    st.subheader("📌 Key Points")
-    for point in analysis.key_points:
-        st.markdown(f"- {point.point}")
-
-    st.subheader("❓ Ask a Question")
-    user_question = st.text_input("What do you want to know?")
-    if user_question:
-        with st.spinner("Searching for an answer..."):
-            answer = answer_question(vectorstore, user_question)
-        st.success(f"💬 Answer: {answer}")
+# --- PDF Upload + Analysis Route ---
+@app.post("/analyze-pdf", response_model=DocumentAnalysis)
+async def analyze_pdf(file: UploadFile = File(...)):
+    text = extract_text_from_pdf(file)
+    analysis = analyze_text_structured(text)
+    return analysis
+
+# --- Question Answering Route ---
+@app.post("/qa", response_model=QAResponse)
+async def ask_question(qa_request: QARequest):
+    answer = answer_question(qa_request.question, qa_request.context)
+    return QAResponse(answer=answer)
+
+# --- Embedding Search (FAISS) Demo ---
+@app.post("/search-chunks")
+async def search_chunks(file: UploadFile = File(...), query: str = ""):
+    text = extract_text_from_pdf(file)
+    chunks = text_splitter.split_text(text)
+    documents = [LangchainDocument(page_content=chunk) for chunk in chunks]
+
+    # Create FAISS vector store
+    db = FAISS.from_documents(documents, embedding_model)
+
+    # Similarity search
+    results = db.similarity_search(query, k=3)
+    return {"results": [doc.page_content for doc in results]}