app.py

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/12.jpg",False,False, "Punjabi"],
        ["Examples/26.jpg",False,False, "Punjabi"],
        ["Examples/36.jpg",False,False, "Punjabi"]],
        # ["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
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
        )

        # Initialize ChromaDB with new configuration
        self.client = chromadb.PersistentClient(path="chroma_db")
        
        # Create or get collection with embeddings
        self.collection = self.client.get_or_create_collection(
            name="my_collection",
            metadata={"hnsw:space": "cosine"}
        )

        # Initialize vectorstore
        self.vectorstore = Chroma(
            client=self.client,
            collection_name="my_collection",
            embedding_function=self.embeddings
        )
        
        self.retriever = self.vectorstore.as_retriever()
        self.ingested_files = []
        
        self.prompt = ChatPromptTemplate.from_messages([...])  # Your existing prompt
        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 []

    def get_example_questions(self):
        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:
    gr.Markdown("# RAG Interface")
    gr.Markdown("Manage documents and ask questions with a Retrieval-Augmented Generation (RAG) system.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### File Management")
            file_input = gr.File(label="Upload File to Ingest")
            add_file_button = gr.Button("Ingest File")
            
            ingested_files_box = gr.Dataframe(
                headers=["Files"], 
                datatype="str",
                row_count=4,
                interactive=False
            )

            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")

            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
            )
            add_file_button.click(
                fn=data_obj.add_file,
                inputs=file_input,
                outputs=ingested_files_box
            )

            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()
                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
            )

        with gr.Column():
            gr.Markdown("### Ask a Question")
            example_questions = gr.Radio(
                choices=data_obj.get_example_questions(),
                label="Example Questions"
            )
            question_input = gr.Textbox(label="Enter your question")
            ask_button = gr.Button("Get Answer")
            answer_output = gr.Textbox(label="Answer", interactive=False)

            def set_example_question(example):
                return gr.update(value=example)

            example_questions.change(
                fn=set_example_question,
                inputs=example_questions,
                outputs=question_input
            )
            ask_button.click(
                fn=answer_question,
                inputs=question_input,
                outputs=answer_output
            )


        # 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)

        #     gr.Markdown("### Ask a Question")
        #     example_questions = gr.Radio(choices=data_obj.get_example_questions(), label="Example Questions")
        #     question_input = gr.Textbox(label="Enter your question")
        #     ask_button = gr.Button("Get Answer")
        #     answer_output = gr.Textbox(label="Answer", interactive=False)

        #     def set_example_question(example):
        #         return gr.update(value=example)

        #     example_questions.change(fn=set_example_question, inputs=example_questions, outputs=question_input)
        #     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(share=True)