import os import spaces import nltk nltk.download('punkt',quiet=True) nltk.download('punkt_tab') from doctr.io import DocumentFile from doctr.models import ocr_predictor import gradio as gr from PIL import Image import base64 from utils import HocrParser from happytransformer import HappyTextToText, TTSettings from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,logging from transformers.integrations import deepspeed import re import torch from lang_list import ( LANGUAGE_NAME_TO_CODE, T2TT_TARGET_LANGUAGE_NAMES, TEXT_SOURCE_LANGUAGE_NAMES, ) logging.set_verbosity_error() DEFAULT_TARGET_LANGUAGE = "English" from transformers import SeamlessM4TForTextToText from transformers import AutoProcessor model = SeamlessM4TForTextToText.from_pretrained("facebook/hf-seamless-m4t-large") processor = AutoProcessor.from_pretrained("facebook/hf-seamless-m4t-large") import pytesseract as pt import cv2 # OCR Predictor initialization OCRpredictor = ocr_predictor(det_arch='db_mobilenet_v3_large', reco_arch='crnn_vgg16_bn', pretrained=True) # Grammar Correction Model initialization happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction") grammar_args = TTSettings(num_beams=5, min_length=1) # Spell Check Model initialization OCRtokenizer = AutoTokenizer.from_pretrained("Bhuvana/t5-base-spellchecker", use_fast=False) OCRmodel = AutoModelForSeq2SeqLM.from_pretrained("Bhuvana/t5-base-spellchecker") # zero = torch.Tensor([0]).cuda() # print(zero.device) def correct_spell(inputs): input_ids = OCRtokenizer.encode(inputs, return_tensors='pt') sample_output = OCRmodel.generate( input_ids, do_sample=True, max_length=512, top_p=0.99, num_return_sequences=1 ) res = OCRtokenizer.decode(sample_output[0], skip_special_tokens=True) return res def process_text_in_chunks(text, process_function, max_chunk_size=256): # Split text into sentences sentences = re.split(r'(?<=[.!?])\s+', text) processed_text = "" for sentence in sentences: # Further split long sentences into smaller chunks chunks = [sentence[i:i + max_chunk_size] for i in range(0, len(sentence), max_chunk_size)] for chunk in chunks: processed_text += process_function(chunk) processed_text += " " # Add space after each processed sentence return processed_text.strip() @spaces.GPU(duration=60) def greet(img, apply_grammar_correction, apply_spell_check,lang_of_input): if (lang_of_input=="Hindi"): res = pt.image_to_string(img,lang='hin') _output_name = "RESULT_OCR.txt" open(_output_name, 'w').write(res) return res, _output_name, None if (lang_of_input=="Punjabi"): res = pt.image_to_string(img,lang='pan') _output_name = "RESULT_OCR.txt" open(_output_name, 'w').write(res) return res, _output_name, None img.save("out.jpg") doc = DocumentFile.from_images("out.jpg") output = OCRpredictor(doc) res = "" for obj in output.pages: for obj1 in obj.blocks: for obj2 in obj1.lines: for obj3 in obj2.words: res += " " + obj3.value res += "\n" res += "\n" # Process in chunks for grammar correction if apply_grammar_correction: res = process_text_in_chunks(res, lambda x: happy_tt.generate_text("grammar: " + x, args=grammar_args).text) # Process in chunks for spell check if apply_spell_check: res = process_text_in_chunks(res, correct_spell) _output_name = "RESULT_OCR.txt" open(_output_name, 'w').write(res) # Convert OCR output to searchable PDF _output_name_pdf="RESULT_OCR.pdf" xml_outputs = output.export_as_xml() parser = HocrParser() base64_encoded_pdfs = list() for i, (xml, img) in enumerate(zip(xml_outputs, doc)): xml_element_tree = xml[1] parser.export_pdfa(_output_name_pdf, hocr=xml_element_tree, image=img) with open(_output_name_pdf, 'rb') as f: base64_encoded_pdfs.append(base64.b64encode(f.read())) return res, _output_name, _output_name_pdf # Gradio Interface for OCR demo_ocr = gr.Interface( fn=greet, inputs=[ gr.Image(type="pil"), gr.Checkbox(label="Apply Grammar Correction"), gr.Checkbox(label="Apply Spell Check"), gr.Dropdown(["English","Hindi","Punjabi"], label="Select Language", value="English") ], outputs=[ gr.Textbox(label="OCR Text"), gr.File(label="Text file"), gr.File(label="Searchable PDF File(English only)") ], title="OCR with Grammar and Spell Check", description="Upload an image to get the OCR results. Optionally, apply grammar and spell check.", examples=[ ["Examples/Book.png",False, False, "English"], ["Examples/News.png",False, False, "English"], ["Examples/Manuscript.jpg",False, False, "English"], ["Examples/Files.jpg",False, False, "English"], ["Examples/Hindi.jpg",False, False, "Hindi"], ["Examples/Hindi-manu.jpg",False, False, "Hindi"], ["Examples/Punjabi_machine.png",False, False, "Punjabi"]], cache_examples=False ) # demo_ocr.launch(debug=True) def split_text_into_batches(text, max_tokens_per_batch): sentences = nltk.sent_tokenize(text) # Tokenize text into sentences batches = [] current_batch = "" for sentence in sentences: if len(current_batch) + len(sentence) + 1 <= max_tokens_per_batch: # Add 1 for space current_batch += sentence + " " # Add sentence to current batch else: batches.append(current_batch.strip()) # Add current batch to batches list current_batch = sentence + " " # Start a new batch with the current sentence if current_batch: batches.append(current_batch.strip()) # Add the last batch return batches @spaces.GPU(duration=60) def run_t2tt(file_uploader , input_text: str, source_language: str, target_language: str) -> (str, bytes): if file_uploader is not None: with open(file_uploader, 'r') as file: input_text=file.read() source_language_code = LANGUAGE_NAME_TO_CODE[source_language] target_language_code = LANGUAGE_NAME_TO_CODE[target_language] max_tokens_per_batch= 2048 batches = split_text_into_batches(input_text, max_tokens_per_batch) translated_text = "" for batch in batches: text_inputs = processor(text=batch, src_lang=source_language_code, return_tensors="pt") output_tokens = model.generate(**text_inputs, tgt_lang=target_language_code) translated_batch = processor.decode(output_tokens[0].tolist(), skip_special_tokens=True) translated_text += translated_batch + " " output=translated_text.strip() _output_name = "result.txt" open(_output_name, 'w').write(output) return str(output), _output_name with gr.Blocks() as demo_t2tt: with gr.Row(): with gr.Column(): with gr.Group(): file_uploader = gr.File(label="Upload a text file (Optional)") input_text = gr.Textbox(label="Input text") with gr.Row(): source_language = gr.Dropdown( label="Source language", choices=TEXT_SOURCE_LANGUAGE_NAMES, value="Punjabi", ) target_language = gr.Dropdown( label="Target language", choices=T2TT_TARGET_LANGUAGE_NAMES, value=DEFAULT_TARGET_LANGUAGE, ) btn = gr.Button("Translate") with gr.Column(): output_text = gr.Textbox(label="Translated text") output_file = gr.File(label="Translated text file") gr.Examples( examples=[ [ None, "The annual harvest festival of Baisakhi in Punjab showcases the region's agricultural prosperity and cultural vibrancy. This joyful occasion brings together people of all ages to celebrate with traditional music, dance, and feasts, reflecting the enduring spirit and community bond of Punjab's people", "English", "Punjabi", ], [ None, "It contains. much useful information about administrative, revenue, judicial and ecclesiastical activities in various areas which, it is hoped, would supplement the information available in official records.", "English", "Hindi", ], [ None, "दुनिया में बहुत सी अलग-अलग भाषाएं हैं और उनमें अपने वर्ण और शब्दों का भंडार होता है. इसमें में कुछ उनके अपने शब्द होते हैं तो कुछ ऐसे भी हैं, जो दूसरी भाषाओं से लिए जाते हैं.", "Hindi", "Punjabi", ], [ None, "ਸੂੂਬੇ ਦੇ ਕਈ ਜ਼ਿਲ੍ਹਿਆਂ ’ਚ ਬੁੱਧਵਾਰ ਸਵੇਰੇ ਸੰਘਣੀ ਧੁੰਦ ਛਾਈ ਰਹੀ ਤੇ ਤੇਜ਼ ਹਵਾਵਾਂ ਨੇ ਕਾਂਬਾ ਹੋਰ ਵਧਾ ਦਿੱਤਾ। ਸੱਤ ਸ਼ਹਿਰਾਂ ’ਚ ਦਿਨ ਦਾ ਤਾਪਮਾਨ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੇ ਆਸਪਾਸ ਰਿਹਾ। ਸੂਬੇ ’ਚ ਵੱਧ ਤੋਂ ਵੱਧ ਤਾਪਮਾਨ ’ਚ ਵੀ ਦਸ ਡਿਗਰੀ ਸੈਲਸੀਅਸ ਦੀ ਗਿਰਾਵਟ ਦਰਜ ਕੀਤੀ ਗਈ", "Punjabi", "English", ], ], inputs=[file_uploader ,input_text, source_language, target_language], outputs=[output_text, output_file], fn=run_t2tt, cache_examples=False, api_name=False, ) gr.on( triggers=[input_text.submit, btn.click], fn=run_t2tt, inputs=[file_uploader, input_text, source_language, target_language], outputs=[output_text, output_file], api_name="t2tt", ) #RAG import utils from langchain_mistralai import ChatMistralAI from langchain_core.prompts import ChatPromptTemplate from langchain_core.output_parsers import StrOutputParser from langchain_community.vectorstores import Chroma from langchain_huggingface import HuggingFaceEmbeddings from langchain_core.runnables import RunnablePassthrough import chromadb.api chromadb.api.client.SharedSystemClient.clear_system_cache() os.environ['MISTRAL_API_KEY'] = 'XuyOObDE7trMbpAeI7OXYr3dnmoWy3L0' class VectorData(): def __init__(self): embedding_model_name = 'l3cube-pune/punjabi-sentence-similarity-sbert' model_kwargs = {'device':'cpu',"trust_remote_code": True} self.embeddings = HuggingFaceEmbeddings( model_name=embedding_model_name, model_kwargs=model_kwargs ) self.vectorstore = Chroma(persist_directory="chroma_db", embedding_function=self.embeddings) self.retriever = self.vectorstore.as_retriever() self.ingested_files = [] self.prompt = ChatPromptTemplate.from_messages( [ ( "system", """Answer the question based on the given context. Dont give any ans if context is not valid to question. Always give the source of context: {context} """, ), ("human", "{question}"), ] ) self.llm = ChatMistralAI(model="mistral-large-latest") self.rag_chain = ( {"context": self.retriever, "question": RunnablePassthrough()} | self.prompt | self.llm | StrOutputParser() ) def add_file(self,file): if file is not None: self.ingested_files.append(file.name.split('/')[-1]) self.retriever, self.vectorstore = utils.add_doc(file,self.vectorstore) self.rag_chain = ( {"context": self.retriever, "question": RunnablePassthrough()} | self.prompt | self.llm | StrOutputParser() ) return [[name] for name in self.ingested_files] def delete_file_by_name(self,file_name): if file_name in self.ingested_files: self.retriever, self.vectorstore = utils.delete_doc(file_name,self.vectorstore) self.ingested_files.remove(file_name) return [[name] for name in self.ingested_files] def delete_all_files(self): self.ingested_files.clear() self.retriever, self.vectorstore = utils.delete_all_doc(self.vectorstore) return [] data_obj = VectorData() # Function to handle question answering def answer_question(question): if question.strip(): return f'{data_obj.rag_chain.invoke(question)}' return "Please enter a question." with gr.Blocks() as rag_interface: # Title and Description gr.Markdown("# RAG Interface") gr.Markdown("Manage documents and ask questions with a Retrieval-Augmented Generation (RAG) system.") with gr.Row(): # Left Column: File Management with gr.Column(): gr.Markdown("### File Management") # File upload and ingest file_input = gr.File(label="Upload File to Ingest") add_file_button = gr.Button("Ingest File") # Scrollable list for ingested files ingested_files_box = gr.Dataframe( headers=["Files"], datatype="str", row_count=4, # Limits the visible rows to create a scrollable view interactive=False ) # Radio buttons to choose delete option delete_option = gr.Radio(choices=["Delete by File Name", "Delete All Files"], label="Delete Option") file_name_input = gr.Textbox(label="Enter File Name to Delete", visible=False) delete_button = gr.Button("Delete Selected") # Show or hide file name input based on delete option selection def toggle_file_input(option): return gr.update(visible=(option == "Delete by File Name")) delete_option.change(fn=toggle_file_input, inputs=delete_option, outputs=file_name_input) # Handle file ingestion add_file_button.click( fn=data_obj.add_file, inputs=file_input, outputs=ingested_files_box ) # Handle delete based on selected option def delete_action(delete_option, file_name): if delete_option == "Delete by File Name" and file_name: return data_obj.delete_file_by_name(file_name) elif delete_option == "Delete All Files": return data_obj.delete_all_files() else: return [[name] for name in data_obj.ingested_files] delete_button.click( fn=delete_action, inputs=[delete_option, file_name_input], outputs=ingested_files_box ) # Right Column: Question Answering with gr.Column(): gr.Markdown("### Ask a Question") # Question input question_input = gr.Textbox(label="Enter your question") # Get answer button and answer output ask_button = gr.Button("Get Answer") answer_output = gr.Textbox(label="Answer", interactive=False) ask_button.click(fn=answer_question, inputs=question_input, outputs=answer_output) with gr.Blocks() as demo: with gr.Tabs(): with gr.Tab(label="OCR"): demo_ocr.render() with gr.Tab(label="Translate"): demo_t2tt.render() with gr.Tab(label="RAG"): rag_interface.render() if __name__ == "__main__": demo.launch()