Inori_Chatbot / pdf_tab.py
SaitoHoujou's picture
Initialize Deployment
0d061c3 verified
import gradio as gr
import requests
from config import backend_url
file_id = None
def upload_file(file):
global file_id
if not file:
return "No file selected"
with open(file.name, "rb") as f:
files ={
"file": (file.name,f, "application/pdf")
}
res= requests.post(f"{backend_url}/api/gemini/upload-pdf", files=files)
if res.status_code == 200:
file_id = res.json().get("file_id")
return "File uploaded."
else:
return "Upload failed!"
def query_pdf(question):
if not file_id:
return "Please upload a PDF first!"
payload = {
"file_id": file_id,
"question" : question
}
with requests.post(
f"{backend_url}/api/gemini/query-pdf",
json=payload,
stream=True
) as res:
if res.status_code != 200:
yield f"Something's not right: {res.status_code} - {res.text}"
return
reply = ""
for chunk in res.iter_content(chunk_size=None, decode_unicode=True):
reply += chunk
yield reply
def create_pdf_tab():
with gr.Tab("PDF"):
pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
upload_status = gr.Textbox(label="Upload Status")
question = gr.Textbox(label="Your question")
answer = gr.Textbox(label="Reply of Chatbot")
btn = gr.Button("Upload PDF")
query_btn = gr.Button("Ask")
btn.click(upload_file, [pdf_file], [upload_status])
query_btn.click(query_pdf, [question], [answer])
question.submit(query_pdf, [question], [answer])