app.py

import gradio as gr
import os
import time

from langchain.document_loaders import OnlinePDFLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.llms import OpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import ConversationalRetrievalChain
from langchain import PromptTemplate
from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
import requests
from PIL import Image
import torch



# _template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question.
# Chat History:
# {chat_history}
# Follow Up Input: {question}
# Standalone question:"""

# CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)

# template = """
# You are given the following extracted parts of a long document and a question. Provide a short structured answer.
# If you don't know the answer, look on the web. Don't try to make up an answer.
# Question: {question}
# =========
# {context}
# =========
# Answer in Markdown:"""

torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/val/png/20294671002019.png', 'chart_example.png')
torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/test/png/multi_col_1081.png', 'chart_example_2.png')
torch.hub.download_url_to_file('https://raw.githubusercontent.com/vis-nlp/ChartQA/main/ChartQA%20Dataset/test/png/18143564004789.png', 'chart_example_3.png')
torch.hub.download_url_to_file('https://sharkcoder.com/files/article/matplotlib-bar-plot.png', 'chart_example_4.png')


model_name = "google/matcha-chartqa"
model = Pix2StructForConditionalGeneration.from_pretrained(model_name)
processor = Pix2StructProcessor.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

def filter_output(output):
    return output.replace("<0x0A>", "")

def chart_qa(image, question):
    inputs = processor(images=image, text=question, return_tensors="pt").to(device)
    predictions = model.generate(**inputs, max_new_tokens=512)
    return filter_output(processor.decode(predictions[0], skip_special_tokens=True))

def loading_pdf():
    return "Loading..."


def pdf_changes(pdf_doc, open_ai_key):
    if open_ai_key is not None:
        os.environ['OPENAI_API_KEY'] = open_ai_key
        loader = OnlinePDFLoader(pdf_doc.name)
        documents = loader.load()
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
        texts = text_splitter.split_documents(documents)
        embeddings = OpenAIEmbeddings()
        db = Chroma.from_documents(texts, embeddings)
        retriever = db.as_retriever()
        global qa
        qa = ConversationalRetrievalChain.from_llm(
            llm=OpenAI(temperature=0.5), 
            retriever=retriever, 
            return_source_documents=True)
        return "Ready"
    else:
        return "You forgot OpenAI API key"

def add_text(history, text):
    history = history + [(text, None)]
    return history, ""

def bot(history):
    response = infer(history[-1][0], history)
    history[-1][1] = ""
    
    for character in response:     
        history[-1][1] += character
        time.sleep(0.05)
        yield history
    

def infer(question, history):  
    res = []
    for human, ai in history[:-1]:
        pair = (human, ai)
        res.append(pair)
    
    chat_history = res
    #print(chat_history)
    query = question
    result = qa({"question": query, "chat_history": chat_history})
    #print(result)
    return result["answer"]

css="""
#col-container {max-width: 700px; margin-left: auto; margin-right: auto;}
"""

title = """
<div style="text-align: center;">
    <h1>YnP LangChain Test </h1>
    <p style="text-align: center;">Please specify OpenAI Key before use</p>
</div>
"""


# with gr.Blocks(css=css) as demo:
#     with gr.Column(elem_id="col-container"):
#         gr.HTML(title)
        
#         with gr.Column():
#             openai_key = gr.Textbox(label="You OpenAI API key", type="password")
#             pdf_doc = gr.File(label="Load a pdf", file_types=['.pdf'], type="file")
#             with gr.Row():
#                 langchain_status = gr.Textbox(label="Status", placeholder="", interactive=False)
#                 load_pdf = gr.Button("Load pdf to langchain")
        
#         chatbot = gr.Chatbot([], elem_id="chatbot").style(height=350)
#         question = gr.Textbox(label="Question", placeholder="Type your question and hit Enter ")
#         submit_btn = gr.Button("Send Message")
        
#     load_pdf.click(loading_pdf, None, langchain_status, queue=False)    
#     load_pdf.click(pdf_changes, inputs=[pdf_doc, openai_key], outputs=[langchain_status], queue=False)
#     question.submit(add_text, [chatbot, question], [chatbot, question]).then(
#         bot, chatbot, chatbot
#     )
#     submit_btn.click(add_text, [chatbot, question], [chatbot, question]).then(
#         bot, chatbot, chatbot)

# demo.launch()


"""functions"""

def load_file():
    return "Loading..."

def load_xlsx(name):
    import pandas as pd
    
    xls_file = rf'{name}'
    data = pd.read_excel(xls_file)  
    return data

def table_loader(table_file, open_ai_key):
    import os
    from langchain.llms import OpenAI
    from langchain.agents import create_pandas_dataframe_agent    
    from pandas import read_csv
    
    global agent
    if open_ai_key is not None:
        os.environ['OPENAI_API_KEY'] = open_ai_key
    else:
        return "Enter API"
    
    if table_file.name.endswith('.xlsx') or table_file.name.endswith('.xls'):
        data = load_xlsx(table_file.name)
        agent = create_pandas_dataframe_agent(OpenAI(temperature=0), data)
        return "Ready!"
    elif table_file.name.endswith('.csv'):
        data = read_csv(table_file.name)
        agent = create_pandas_dataframe_agent(OpenAI(temperature=0), data)
        return "Ready!"
    else:
        return "Wrong file format! Upload excel file or csv!"
    
def run(query):
    from langchain.callbacks import get_openai_callback
    
    with get_openai_callback() as cb:
        response = (agent.run(query))
        costs = (f"Total Cost (USD): ${cb.total_cost}")
        output = f'{response} \n {costs}'
        return output
    
def respond(message, chat_history):
    import time
    
    bot_message = run(message)
    chat_history.append((message, bot_message))
    time.sleep(0.5)
    return "", chat_history


with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.HTML(title)
        key = gr.Textbox(
                show_label=False,
                placeholder="Your OpenAI key",
                type = 'password',
                ).style(container=False)

    # PDF processing tab
    with gr.Tab("PDFs"):
        
        with gr.Row(): 
            
            with gr.Column(scale=0.5):
                langchain_status = gr.Textbox(label="Status", placeholder="", interactive=False)
                load_pdf = gr.Button("Load pdf to langchain")
                
            with gr.Column(scale=0.5):
                pdf_doc = gr.File(label="Load a pdf", file_types=['.pdf'], type="file")
                
                
        with gr.Row():
            
            with gr.Column(scale=1): 
                chatbot = gr.Chatbot([], elem_id="chatbot").style(height=350)
                
        with gr.Row():
            
            with gr.Column(scale=0.85):
                question = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter, or upload an image",
                ).style(container=False)
                
            with gr.Column(scale=0.15, min_width=0):
                clr_btn = gr.Button("Clear!")
                
    load_pdf.click(loading_pdf, None, langchain_status, queue=False)    
    load_pdf.click(pdf_changes, inputs=[pdf_doc, key], outputs=[langchain_status], queue=True)
    question.submit(add_text, [chatbot, question], [chatbot, question]).then(
        bot, chatbot, chatbot
    )
                
    # XLSX and CSV processing tab
    with gr.Tab("Spreadsheets"):
        with gr.Row(): 
            
            with gr.Column(scale=0.5):
                status_sh = gr.Textbox(label="Status", placeholder="", interactive=False)
                load_table = gr.Button("Load csv|xlsx to langchain")
                
            with gr.Column(scale=0.5):
                raw_table = gr.File(label="Load a table file (xls or csv)", file_types=['.csv, xlsx, xls'], type="file")
                
                
        with gr.Row():
            
            with gr.Column(scale=1): 
                chatbot_sh = gr.Chatbot([], elem_id="chatbot").style(height=350)
                
                
        with gr.Row():
            
            with gr.Column(scale=0.85):
                question_sh = gr.Textbox(
                show_label=False,
                placeholder="Enter text and press enter, or upload an image",
                ).style(container=False)
                
            with gr.Column(scale=0.15, min_width=0):
                clr_btn = gr.Button("Clear!")
            
    load_table.click(load_file, None, status_sh, queue=False)    
    load_table.click(table_loader, inputs=[raw_table, key], outputs=[status_sh], queue=False)
   
    question_sh.submit(respond, [question_sh, chatbot_sh], [question_sh, chatbot_sh])
    clr_btn.click(lambda: None, None, chatbot_sh, queue=False)
    

    with gr.Tab("Charts"):
            image = gr.Image(type="pil", label="Chart")
            question = gr.Textbox(label="Question")
            load_chart = gr.Button("Load chart and question!")
            answer = gr.Textbox(label="Model Output")
            
    load_chart.click(chart_qa, [image, question], answer)

        
demo.queue(concurrency_count=3)
demo.launch()