Gradio / app.py
ajalisatgi's picture
Update app.py
4d2d551 verified
raw
history blame
3.62 kB
import gradio as gr
import openai
import os
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.schema import Document
from sentence_transformers import SentenceTransformer
from datasets import load_dataset
import nltk
# βœ… Load the Sentence Transformer Embedding Model
model_name = "sentence-transformers/all-MiniLM-L6-v2"
embedding_model = HuggingFaceEmbeddings(model_name=model_name)
# βœ… Set OpenAI API Key
openai.api_key = os.getenv("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
")
# βœ… Download NLTK Tokenizer
nltk.download('punkt')
# βœ… Load and Chunk Dataset
def chunk_documents(documents, max_chunk_size=500):
chunks = []
for doc in documents:
sentences = nltk.sent_tokenize(doc)
current_chunk = ""
for sentence in sentences:
if len(current_chunk) + len(sentence) <= max_chunk_size:
current_chunk += sentence + " "
else:
chunks.append(current_chunk.strip())
current_chunk = sentence + " "
if current_chunk:
chunks.append(current_chunk.strip())
return chunks
# βœ… Load Dataset and Prepare ChromaDB
dataset = load_dataset("rungalileo/ragbench", "techqa") # Example dataset
original_documents = dataset['train']['documents']
chunked_documents = chunk_documents(original_documents)
persist_directory = "chroma_db_directory"
documents = [Document(page_content=chunk) for chunk in chunked_documents]
# βœ… Initialize ChromaDB
vectordb = Chroma.from_documents(
documents=documents,
embedding=embedding_model,
persist_directory=persist_directory
)
vectordb.persist()
# βœ… Function to Retrieve Relevant Documents
def retrieve_documents(question, k=5):
docs = vectordb.similarity_search(question, k=k)
if not docs:
return ["⚠️ No relevant documents found. Try a different query."]
return [doc.page_content for doc in docs]
# βœ… Function to Generate AI Response
def generate_response(question, context):
if not context or "No relevant documents found." in context:
return "No relevant context available. Try a different query."
full_prompt = f"Context: {context}\n\nQuestion: {question}"
try:
response = openai.ChatCompletion.create(
model="gpt-4",
messages=[
{"role": "system", "content": "You are an AI assistant that answers user queries based on the given context."},
{"role": "user", "content": full_prompt}
],
max_tokens=300,
temperature=0.7
)
return response['choices'][0]['message']['content'].strip()
except Exception as e:
return f"Error generating response: {str(e)}"
# βœ… Full RAG Pipeline
def rag_pipeline(question):
retrieved_docs = retrieve_documents(question, k=5)
context = " ".join(retrieved_docs)
response = generate_response(question, context)
return response, "\n\n".join(retrieved_docs)
# βœ… Gradio UI Interface
iface = gr.Interface(
fn=rag_pipeline,
inputs=gr.Textbox(label="Enter your question"),
outputs=[
gr.Textbox(label="Generated Response"),
gr.Textbox(label="Retrieved Documents")
],
title="RAG-Based Question Answering System",
description="Enter a question and retrieve relevant documents with AI-generated response."
)
# βœ… Launch the Gradio App
if __name__ == "__main__":
iface.launch()