text-rag / app.py
saravan5757's picture
Create app.py
d0bda0e verified
Raw
History Blame Contribute Delete
2.77 kB
import gradio as gr
import faiss
import numpy as np
from sentence_transformers import SentenceTransformer
from transformers import pipeline
from PyPDF2 import PdfReader
# Load AI models
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2") # Converts text to numbers
llm_pipeline = pipeline("text2text-generation", model="google/flan-t5-small") # AI that answers questions
# Memory (Database) to store text
index = None
chunks = []
# Load and process document
def load_document(file):
global index, chunks
text = ""
# Determine file path or object
# If file is a string, it's a file path
if isinstance(file, str):
file_path = file
else:
# If file is not a string, try to use its .name attribute
file_path = file.name
# Read PDF or text file
if file_path.endswith(".pdf"):
reader = PdfReader(file_path)
text = "\n".join(
[page.extract_text() for page in reader.pages if page.extract_text()]
)
else:
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
# Break text into small parts
sentences = text.split(". ")
chunks = [" ".join(sentences[i:i + 5]) for i in range(0, len(sentences), 5)]
# Create embeddings and store in FAISS
embeddings = np.array([embedding_model.encode(chunk) for chunk in chunks])
index = faiss.IndexFlatL2(embeddings.shape[1])
index.add(embeddings)
return "πŸ“„ Document is ready! Now ask your question."
# Find and answer questions
def get_answer(query):
if index is None:
return "❌ Please upload a document first."
# Find best matching text
query_embedding = embedding_model.encode(query).reshape(1, -1)
distances, indices = index.search(query_embedding, 3)
retrieved_text = " ".join([chunks[i] for i in indices[0]])
# Ask AI to generate an answer
input_text = f"Question: {query}\nContext: {retrieved_text}"
response = llm_pipeline(input_text, max_length=100)[0]['generated_text']
return response
# Webpage design
with gr.Blocks() as demo:
gr.Markdown("# πŸ“š Smart Study Helper")
file_input = gr.File(label="Upload a textbook or notes (PDF/TXT)", file_types=[".pdf", ".txt"])
upload_button = gr.Button("Process Document")
status_text = gr.Textbox(label="πŸ“’ Status", interactive=False)
query_input = gr.Textbox(label="Ask a question from the document:")
query_button = gr.Button("Get Answer")
output_text = gr.Textbox(label="πŸ€– AI Answer", interactive=False)
upload_button.click(load_document, inputs=file_input, outputs=status_text)
query_button.click(get_answer, inputs=query_input, outputs=output_text)
# Run the app
demo.launch()