Spaces:
Build error
Build error
File size: 3,265 Bytes
b6b32c0 |
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 |
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
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() |