import requests import json import gradio as gr import pdfplumber import pandas as pd from datetime import datetime from google.oauth2.service_account import Credentials from cnocr import CnOcr import gspread from sentence_transformers import SentenceTransformer, models, util # Load credentials for Google Sheets credentials = Credentials.from_service_account_file("credentials.json", scopes=["https://www.googleapis.com/auth/spreadsheets"]) client = gspread.authorize(credentials) sheet = client.open_by_url("https://docs.google.com/spreadsheets/d/16H4M-8hHdOhI68vDIsDFT6T2xcGEvm0A7o5uFlmrzrQ/edit?usp=sharing").sheet1 # Initialize models and utilities word_embedding_model = models.Transformer('sentence-transformers/all-MiniLM-L6-v2', do_lower_case=True) pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension(), pooling_mode='cls') embedder = SentenceTransformer(modules=[word_embedding_model, pooling_model]) ocr = CnOcr() # API URLs and headers chat_url = 'https://Raghav001-API.hf.space/chatpdf' chat_emd = 'https://Raghav001-API.hf.space/embedd' headers = {'Content-Type': 'application/json'} # Global variables history_max_len = 500 all_max_len = 3000 bot = [] def record_to_sheet(timestamp, user_input, answer): row = [timestamp, user_input, answer] sheet.append_row(row) def doc_emb(doc): texts = doc.split('\n') emb_list = embedder.encode(texts) print('emb_list', emb_list) print('\n'.join(texts)) gr.Textbox.update(value="") return texts, emb_list, gr.Textbox.update(visible=True), gr.Button.update(visible=True), gr.Markdown.update(value="""success! Let's talk"""), gr.Chatbot.update(visible=True) def get_response(msg, bot, doc_text_list, doc_embeddings): gr.Textbox.update(value="") now_len = len(msg) req_json = {'question': msg} his_bg = -1 for i in range(len(bot) - 1, -1, -1): if now_len + len(bot[i][0]) + len(bot[i][1]) > history_max_len: break now_len += len(bot[i][0]) + len(bot[i][1]) his_bg = i req_json['history'] = [] if his_bg == -1 else bot[his_bg:] query_embedding = embedder.encode([msg]) cos_scores = util.cos_sim(query_embedding, doc_embeddings)[0] score_index = [[score, index] for score, index in zip(cos_scores, [i for i in range(len(cos_scores))])] score_index.sort(key=lambda x: x[0], reverse=True) print('score_index:\n', score_index) print('doc_emb_state', doc_emb_state) index_set, sub_doc_list = set(), [] for s_i in score_index: doc = doc_text_list[s_i[1]] if now_len + len(doc) > all_max_len: break index_set.add(s_i[1]) now_len += len(doc) # Maybe the paragraph is truncated wrong, so add the upper and lower paragraphs if s_i[1] > 0 and s_i[1] - 1 not in index_set: doc = doc_text_list[s_i[1]-1] if now_len + len(doc) > all_max_len: break index_set.add(s_i[1]-1) now_len += len(doc) if s_i[1] + 1 < len(doc_text_list) and s_i[1] + 1 not in index_set: doc = doc_text_list[s_i[1]+1] if now_len + len(doc) > all_max_len: break index_set.add(s_i[1]+1) now_len += len(doc) index_list = list(index_set) index_list.sort() for i in index_list: sub_doc_list.append(doc_text_list[i]) req_json['doc'] = '' if len(sub_doc_list) == 0 else '\n'.join(sub_doc_list) data = {"content": json.dumps(req_json)} print('data:\n', req_json) result = requests.post(url=chat_url, data=json.dumps(data), headers=headers) res = result.json()['content'] bot.append([msg, res]) record_to_sheet(datetime.now().strftime("%Y-%m-%d %H:%M:%S"), msg, res) return bot[max(0, len(bot) - 3):] def up_file(fls): doc_text_list = [] names = [str(i.name) for i in fls] pdf = [] docs = [] pptx = [] for i in names: if i[-3:] == "pdf": pdf.append(i) elif i[-4:] == "docx": docs.append(i) else: pptx.append(i) # Pdf Extraction for idx, file in enumerate(pdf): with pdfplumber.open(file) as pdf: for i in range(len(pdf.pages)): page = pdf.pages[i] res_list = page.extract_text().split('\n')[:-1] for j in range(len(page.images)): img = page.images[j] file_name = f"{str(time.time())}-{str(i)}-{str(j)}.png" with open(file_name, mode='wb') as f: f.write(img['stream'].get_data()) try: res = ocr.ocr(file_name) except Exception as e: res = [] if len(res) > 0: res_list.append(' '.join([re['text'] for re in res])) tables = page.extract_tables() for table in tables: df = pd.DataFrame(table[1:], columns=table[0]) try: records = json.loads(df.to_json(orient="records")) for rec in records: res_list.append(json.dumps(rec)) except Exception as e: res_list.append(str(df)) doc_text_list += res_list # PPTX Extraction for i in pptx: loader = UnstructuredPowerPointLoader(i) data = loader.load() doc_text_list.append(data) # Doc Extraction for i in docs: loader = UnstructuredWordDocumentLoader(i) data = loader.load() doc_text_list.append(data) doc_text_list = [str(text).strip() for text in doc_text_list if len(str(text).strip()) > 0] return gr.Textbox.update(value='\n'.join(doc_text_list), visible=True), gr.Button.update(visible=True), gr.Markdown.update(value="Processing") def launch_interface(): with gr.Interface( fn=up_file, inputs="file", outputs=["text", "button", "markdown"], title="Document Chatbot", description="Upload a PDF contract to chat with the AI lawyer." ) as interface: interface.launch() if __name__ == "__main__": launch_interface()