DrishtiSharma commited on
Commit
ad3d855
Β·
verified Β·
1 Parent(s): 07de8c7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +96 -0
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.")