Spaces:
Sleeping
Sleeping
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() | |