fedor-ch's picture
charts tab (#4)
852b083
raw
history blame contribute delete
No virus
10 kB
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()