Spaces:
Build error
Build error
Update app.py
Browse files
app.py
CHANGED
@@ -0,0 +1,96 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
from langchain_groq import ChatGroq
|
3 |
+
from langchain_ollama import ChatOllama
|
4 |
+
from langchain_community.document_loaders import PyMuPDFLoader
|
5 |
+
from langchain.schema import Document
|
6 |
+
from langchain_core.prompts import PromptTemplate
|
7 |
+
from langchain_core.output_parsers import JsonOutputParser
|
8 |
+
from pydantic import BaseModel, Field
|
9 |
+
import os
|
10 |
+
import json
|
11 |
+
from prompt import REAG_SYSTEM_PROMPT, rag_prompt # Import prompts
|
12 |
+
|
13 |
+
# Hugging Face Spaces API Key Handling
|
14 |
+
st.set_page_config(page_title="ReAG", layout="centered")
|
15 |
+
|
16 |
+
# Set API Key using Hugging Face Secrets
|
17 |
+
os.environ["GROQ_API_KEY"] = st.secrets["GROQ_API_KEY"]
|
18 |
+
|
19 |
+
# Initialize LLM models
|
20 |
+
llm_relevancy = ChatGroq(model="llama-3.3-70b-versatile", temperature=0)
|
21 |
+
llm = ChatOllama(model="deepseek-r1:14b", temperature=0.6, max_tokens=3000)
|
22 |
+
|
23 |
+
# Define schema for extracted content
|
24 |
+
class ResponseSchema(BaseModel):
|
25 |
+
content: str = Field(..., description="Relevant content from the document")
|
26 |
+
reasoning: str = Field(..., description="Why this content was selected")
|
27 |
+
is_irrelevant: bool = Field(..., description="True if the content is irrelevant")
|
28 |
+
|
29 |
+
class RelevancySchemaMessage(BaseModel):
|
30 |
+
source: ResponseSchema
|
31 |
+
|
32 |
+
relevancy_parser = JsonOutputParser(pydantic_object=RelevancySchemaMessage)
|
33 |
+
|
34 |
+
# Function to format document
|
35 |
+
def format_doc(doc: Document) -> str:
|
36 |
+
return f"Document_Title: {doc.metadata.get('title', 'Unknown')}\nPage: {doc.metadata.get('page', 'Unknown')}\nContent: {doc.page_content}"
|
37 |
+
|
38 |
+
# Extract relevant context function
|
39 |
+
def extract_relevant_context(question, documents):
|
40 |
+
result = []
|
41 |
+
with st.spinner("π Extracting relevant content from document..."):
|
42 |
+
for doc in documents:
|
43 |
+
formatted_documents = format_doc(doc)
|
44 |
+
system_prompt = f"{REAG_SYSTEM_PROMPT}\n\n# Available source\n\n{formatted_documents}"
|
45 |
+
prompt = f"""Determine if the 'Available source' content is sufficient to answer the QUESTION.
|
46 |
+
QUESTION: {question}
|
47 |
+
RESPONSE FORMAT (Strict JSON):
|
48 |
+
```json
|
49 |
+
{{
|
50 |
+
"content": "Extracted relevant content",
|
51 |
+
"reasoning": "Why this was chosen",
|
52 |
+
"is_irrelevant": false
|
53 |
+
}}
|
54 |
+
```
|
55 |
+
"""
|
56 |
+
messages = [{"role": "system", "content": system_prompt},
|
57 |
+
{"role": "user", "content": prompt}]
|
58 |
+
response = llm_relevancy.invoke(messages)
|
59 |
+
formatted_response = relevancy_parser.parse(response.content)
|
60 |
+
result.append(formatted_response)
|
61 |
+
|
62 |
+
final_context = [item['content'] for item in result if not item['is_irrelevant']]
|
63 |
+
return final_context
|
64 |
+
|
65 |
+
# Generate response using RAG Prompt
|
66 |
+
def generate_response(question, final_context):
|
67 |
+
with st.spinner("π Generating AI response..."):
|
68 |
+
prompt = PromptTemplate(template=rag_prompt, input_variables=["question", "context"])
|
69 |
+
chain = prompt | llm
|
70 |
+
response = chain.invoke({"question": question, "context": final_context})
|
71 |
+
return response.content.split("\n\n")[-1]
|
72 |
+
|
73 |
+
# Streamlit UI
|
74 |
+
st.title("π RAG-based Knowledge Retrieval App")
|
75 |
+
st.markdown("Upload a PDF and ask questions based on its content.")
|
76 |
+
|
77 |
+
uploaded_file = st.file_uploader("π Upload PDF", type=["pdf"])
|
78 |
+
if uploaded_file:
|
79 |
+
with st.spinner("π₯ Uploading and processing PDF..."):
|
80 |
+
file_path = f"/tmp/{uploaded_file.name}"
|
81 |
+
with open(file_path, "wb") as f:
|
82 |
+
f.write(uploaded_file.getbuffer())
|
83 |
+
|
84 |
+
loader = PyMuPDFLoader(file_path)
|
85 |
+
docs = loader.load()
|
86 |
+
st.success("β
PDF uploaded and processed successfully!")
|
87 |
+
|
88 |
+
question = st.text_input("β Ask a question about the document:")
|
89 |
+
if st.button("π Get Answer"):
|
90 |
+
if question:
|
91 |
+
final_context = extract_relevant_context(question, docs)
|
92 |
+
if final_context:
|
93 |
+
answer = generate_response(question, final_context)
|
94 |
+
st.success(f"π§ **Response:**\n\n{answer}")
|
95 |
+
else:
|
96 |
+
st.warning("β οΈ No relevant information found in the document.")
|