|
|
import gradio as gr |
|
|
from PyPDF2 import PdfReader |
|
|
import docx |
|
|
|
|
|
def extract_text_from_uploaded_file(file_path): |
|
|
if file_path.endswith(".pdf"): |
|
|
try: |
|
|
reader = PdfReader(file_path) |
|
|
return "\n".join(page.extract_text() or "" for page in reader.pages) |
|
|
except Exception as e: |
|
|
return f"[Error extracting PDF text: {e}]" |
|
|
elif file_path.endswith(".txt"): |
|
|
try: |
|
|
with open(file_path, "r", encoding="utf-8") as f: |
|
|
return f.read() |
|
|
except Exception as e: |
|
|
return f"[Error reading text file: {e}]" |
|
|
elif file_path.endswith(".docx"): |
|
|
try: |
|
|
doc = docx.Document(file_path) |
|
|
return "\n".join([p.text for p in doc.paragraphs]) |
|
|
except Exception as e: |
|
|
return f"[Error reading DOCX: {e}]" |
|
|
else: |
|
|
return "[Unsupported file type for text extraction]" |
|
|
|
|
|
def create_ui(agent): |
|
|
with gr.Blocks(theme=gr.themes.Soft()) as demo: |
|
|
gr.Markdown("<h1 style='text-align: center;'>💊 TxAgent: Therapeutic Reasoning</h1>") |
|
|
chatbot = gr.Chatbot(label="TxAgent", height=600, type="messages") |
|
|
|
|
|
file_upload = gr.File(label="Upload Medical File", file_types=[".pdf", ".txt", ".docx"]) |
|
|
message_input = gr.Textbox(placeholder="Ask a biomedical question...", show_label=False) |
|
|
send_button = gr.Button("Send", variant="primary") |
|
|
conversation_state = gr.State([]) |
|
|
|
|
|
def handle_chat(message, history, conversation, uploaded_file): |
|
|
file_text = "" |
|
|
if uploaded_file: |
|
|
file_text = extract_text_from_uploaded_file(uploaded_file.name) |
|
|
|
|
|
|
|
|
if file_text: |
|
|
message += f"\n\n[Document Content Extracted from Upload]:\n{file_text}" |
|
|
|
|
|
generator = agent.run_gradio_chat( |
|
|
message=message, |
|
|
history=history, |
|
|
temperature=0.3, |
|
|
max_new_tokens=1024, |
|
|
max_token=8192, |
|
|
call_agent=False, |
|
|
conversation=conversation, |
|
|
max_round=30 |
|
|
) |
|
|
for update in generator: |
|
|
yield update |
|
|
|
|
|
inputs = [message_input, chatbot, conversation_state, file_upload] |
|
|
send_button.click(fn=handle_chat, inputs=inputs, outputs=chatbot) |
|
|
message_input.submit(fn=handle_chat, inputs=inputs, outputs=chatbot) |
|
|
|
|
|
gr.Examples( |
|
|
examples=[ |
|
|
["Upload the files"], |
|
|
], |
|
|
inputs=message_input, |
|
|
) |
|
|
|
|
|
return demo |
|
|
|
