Spaces:
Sleeping
Sleeping
File size: 3,413 Bytes
0fbe4cf 5673d2d 0fbe4cf |
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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import gradio as gr
import os
from tqdm import tqdm
from unstructured.partition.pdf import partition_pdf
from langchain.schema.document import Document
import google.generativeai as genai
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
import shutil
from wasabi import msg
from PIL import Image
os.system('sudo apt update')
os.system('sudo apt upgrade')
os.system('sudo apt install poppler-utils')
os.system('sudo apt install tesseract-ocr')
genai.configure(api_key = 'AIzaSyB342Fh-nkRaO38BshbyI4-s0T9orVpsMw')
model = genai.GenerativeModel('gemini-1.5-flash')
files = os.listdir('uploads')
files = [f'uploads/{file}' for file in files]
documents = []
for file in tqdm(files , total = len(files) , leave = False) :
elements = partition_pdf(
filename = file ,
extract_images_in_pdf = True ,
infer_table_structure = True ,
chunking_strategy = 'by_title' ,
max_characters = 4000 ,
new_after_n_chars = 3800 ,
combine_text_under_n_chars = 2000 ,
extract_image_block_output_dir = 'outputs'
)
for element in elements :
element = element.to_dict()
metadata = element['metadata']
if 'text_as_html' in metadata : documents.append(
Document(
page_content = metadata['text_as_html'] ,
metadata = {
'type' : 'text' ,
'metadata' : element
}
)
)
else : documents.append(
Document(
page_content = element['text'] ,
metadata = {
'type' : 'text' ,
'metadata' : element
}
)
)
images = os.listdir('outputs')
images = [f'outputs/{image}' for image in images]
for image in tqdm(images , total = len(images) , leave = False) :
image = Image.open(image)
try :
response = model.generate_content([
image ,
'Explain the Image'
])
response = response.text
except Exception as e : msg.fail(f'----| FAIL : COULDNT CALL THE IMAGE DESCRIPTION API : {e}') ; response = 'COuldnt Call Model for this'
documents.append(
Document(
page_content = response ,
metadata = {
'type' : 'image' ,
'metadata' : {
'image' : image
}
}
)
)
shutil.rmtree('uploads')
vc = FAISS.from_documents(
documents = documents ,
embedding = HuggingFaceEmbeddings(model_name = 'all-MiniLM-L6-v2')
)
def run_rag(query) :
similar_docs = vc.similarity_search(query , k = 4)
context = [doc.page_content for doc in similar_docs]
prompt = f'''
You are a Helpfull Chatbot that helps users with their queries
- You will be provided with a query
- You will be provided with a context as well
Your task is to generate a response to the query based on the context provided
Context : {context}
Query : {query}
'''
response = model.generate_content(prompt)
return response.text
demo = gr.Interface(
fn = run_rag ,
inputs = 'text' ,
outputs = 'text'
)
demo.launch() |