File size: 3,238 Bytes
bad3068 |
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 |
import os
import gradio as gr
import google.generativeai as genai
from typing import List, Tuple
import fitz # PyMuPDF
from sentence_transformers import SentenceTransformer
import numpy as np
import faiss
# Initialize Google API Key
GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
genai.configure(api_key=GOOGLE_API_KEY)
# Placeholder for the app's state
class MyApp:
def __init__(self) -> None:
self.documents = []
self.embeddings = None
self.index = None
self.load_pdf("THEDIA1.pdf")
self.build_vector_db()
def load_pdf(self, file_path: str) -> None:
"""Extracts text from a PDF file and stores it in the app's documents."""
doc = fitz.open(file_path)
self.documents = []
for page_num in range(len(doc)):
page = doc[page_num]
text = page.get_text()
self.documents.append({"page": page_num + 1, "content": text})
print("PDF processed successfully!")
def build_vector_db(self) -> None:
"""Builds a vector database using FAISS and SentenceTransformer embeddings."""
model = SentenceTransformer("all-MiniLM-L6-v2")
embeddings = model.encode([doc["content"] for doc in self.documents])
self.embeddings = np.array(embeddings, dtype="float32")
self.index = faiss.IndexFlatL2(self.embeddings.shape[1])
self.index.add(self.embeddings)
print("Vector database built successfully!")
def search(self, query: str, top_k: int = 5) -> List[Tuple[int, str]]:
"""Searches for the most similar documents based on the query."""
query_embedding = SentenceTransformer("all-MiniLM-L6-v2").encode([query])
distances, indices = self.index.search(np.array(query_embedding, dtype="float32"), top_k)
return [(self.documents[idx]["page"], self.documents[idx]["content"]) for idx in indices[0]]
def generate_response(self, query: str) -> str:
"""Generates a response using the Gemini model based on the query."""
if not GOOGLE_API_KEY:
raise ValueError("GOOGLE_API_KEY is not set. Please set it up.")
generation_config = genai.types.GenerationConfig(
temperature=0.7,
max_output_tokens=512
)
model_name = "gemini-1.5-pro-latest"
model = genai.GenerativeModel(model_name)
response = model.generate_content([query], generation_config=generation_config)
return response[0].text if response else "No response generated."
# Gradio UI setup for interaction
def main():
app = MyApp()
def handle_query(query):
search_results = app.search(query)
response = app.generate_response(query)
return {"Search Results": search_results, "Response": response}
gr.Interface(
fn=handle_query,
inputs=gr.Textbox(placeholder="Enter your query here"),
outputs=[
gr.JSON(label="Search Results"),
gr.Textbox(label="Generated Response")
],
title="Dialectical Behavioral Exercise with Gemini",
description="This app uses Google Gemini to generate responses based on document content."
).launch()
if __name__ == "__main__":
main()
|