Spaces:
Sleeping
Sleeping
File size: 4,114 Bytes
a2aaac9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | import os
import gradio as gr
from groq import Groq
from fpdf import FPDF
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader, TextLoader
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
# -------------------------
# GROQ CLIENT
# -------------------------
client = Groq(api_key=os.environ.get("RAG_API_KEY"))
# -------------------------
# GLOBAL STORAGE
# -------------------------
vectorstore = None
chat_history = []
# -------------------------
# LOAD DOCUMENTS
# -------------------------
def load_documents(files):
documents = []
for file in files:
if file.name.endswith(".pdf"):
loader = PyPDFLoader(file.name)
else:
loader = TextLoader(file.name)
documents.extend(loader.load())
return documents
# -------------------------
# BUILD VECTOR STORE
# -------------------------
def build_vectorstore(files):
global vectorstore
docs = load_documents(files)
splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=100
)
chunks = splitter.split_documents(docs)
embeddings = HuggingFaceEmbeddings(
model_name="sentence-transformers/all-MiniLM-L6-v2"
)
vectorstore = FAISS.from_documents(chunks, embeddings)
return "β
Documents processed successfully."
# -------------------------
# ASK QUESTION
# -------------------------
def ask_question(question):
global chat_history, vectorstore
if vectorstore is None:
return "β Please upload documents first.", ""
docs = vectorstore.similarity_search(question, k=4)
context = "\n\n".join([d.page_content for d in docs])
prompt = f"""
You are a helpful AI assistant.
First, use the document context below to answer.
If the document does not fully answer the question,
you may add relevant general knowledge.
Document Context:
{context}
Question:
{question}
"""
response = client.chat.completions.create(
model="llama-3.3-70b-versatile",
messages=[{"role": "user", "content": prompt}],
)
answer = response.choices[0].message.content
chat_history.append((question, answer))
sources = "\n\n".join(
[f"Source {i+1}:\n{d.page_content[:300]}..." for i, d in enumerate(docs)]
)
return answer, sources
# -------------------------
# EXPORT CHAT TO PDF
# -------------------------
def export_chat_pdf():
pdf = FPDF()
pdf.add_page()
pdf.set_font("Arial", size=12)
pdf.cell(0, 10, "AneesLLM - Chat Export", ln=True)
pdf.ln(5)
for q, a in chat_history:
pdf.multi_cell(0, 8, f"Q: {q}")
pdf.multi_cell(0, 8, f"A: {a}")
pdf.ln(4)
path = "/tmp/AneesLLM_Chat.pdf"
pdf.output(path)
return path
# -------------------------
# GRADIO UI
# -------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
gr.Markdown(
"""
# π€ AneesLLM β RAG Based AI Assistant
Upload documents and ask intelligent questions with sources.
"""
)
with gr.Row():
file_input = gr.File(
file_types=[".pdf", ".txt"],
file_count="multiple",
label="π Upload Documents"
)
process_btn = gr.Button("π Process Documents")
status = gr.Textbox(label="Status")
process_btn.click(
build_vectorstore,
inputs=file_input,
outputs=status
)
gr.Markdown("## π¬ Ask a Question")
question = gr.Textbox(
placeholder="Ask something about your documents...",
label="Your Question"
)
ask_btn = gr.Button("Ask")
answer = gr.Textbox(label="Answer", lines=6)
sources = gr.Textbox(label="Sources Used", lines=6)
ask_btn.click(
ask_question,
inputs=question,
outputs=[answer, sources]
)
export_btn = gr.Button("π Export Chat as PDF")
pdf_file = gr.File(label="Download PDF")
export_btn.click(export_chat_pdf, outputs=pdf_file)
demo.launch() |