ritampatra's picture
Create app.py
5f0ff3b verified
import gradio as gr
from llama_index import SimpleDirectoryReader, GPTSimpleVectorIndex, ServiceContext
from transformers import AutoTokenizer, AutoModelForCausalLM
from PyPDF2 import PdfReader
# Load tokenizer and model from HuggingFace (StableLM)
tokenizer = AutoTokenizer.from_pretrained("StabilityAI/stablelm-tuned-alpha-3b")
model = AutoModelForCausalLM.from_pretrained("StabilityAI/stablelm-tuned-alpha-3b")
# Create service context for the LLM
service_context = ServiceContext.from_defaults(
llm_predictor=(model, tokenizer), # Attach the model and tokenizer
chunk_size=1024
)
# Function to load PDF
def load_pdf(file):
reader = PdfReader(file.name)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
# Function to create an index and query it
def chat_with_pdf(pdf, query):
# Read the PDF content
pdf_text = load_pdf(pdf)
# Use llama-index to create a document index
documents = [pdf_text]
index = GPTSimpleVectorIndex.from_documents(documents, service_context=service_context)
# Query the index
response = index.query(query)
return response.response
# Gradio interface
def chatbot(pdf, query):
if not pdf or not query:
return "Please upload a PDF and enter a query."
response = chat_with_pdf(pdf, query)
return response
# Define Gradio inputs and interface
pdf_input = gr.inputs.File(label="Upload your PDF")
query_input = gr.inputs.Textbox(label="Ask a question about the PDF")
output = gr.outputs.Textbox(label="Chatbot Response")
gr.Interface(fn=chatbot, inputs=[pdf_input, query_input], outputs=output, title="PDF Chatbot").launch()