File size: 2,322 Bytes
68c702d 9baafa1 68c702d 9baafa1 68c702d 4fea39c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import fitz # PyMuPDF
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import gradio as gr
def extract_text_from_pdf(file_path):
doc = fitz.open(file_path)
text = ""
for page in doc:
text += page.get_text()
return text
def analyze_document(file, prompt):
# Determine file type and extract text
if file.name.endswith('.pdf'):
text = extract_text_from_pdf(file.name)
elif file.name.endswith('.txt'):
text = file.read().decode('utf-8')
else:
return "Unsupported file type. Please upload a PDF or TXT file."
# Load the model and tokenizer
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained("THUDM/glm-4v-9b", trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
"THUDM/glm-4v-9b",
torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
low_cpu_mem_usage=True,
trust_remote_code=True
).to(device).eval()
# Prepare inputs
inputs = tokenizer.apply_chat_template([{"role": "user", "content": text + "\n\n" + prompt}],
add_generation_prompt=True, tokenize=True, return_tensors="pt",
return_dict=True)
inputs = inputs.to(device)
# Generate response
gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1}
with torch.no_grad():
outputs = model.generate(**inputs, **gen_kwargs)
outputs = outputs[:, inputs['input_ids'].shape[1]:]
return tokenizer.decode(outputs[0])
# Gradio interface
file_input = gr.File(label="Upload TXT or PDF Document", file_count="single")
prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your structured prompt here")
output_text = gr.Textbox(label="Analysis Result")
iface = gr.Interface(
fn=analyze_document,
# inputs=[
# gr.inputs.File(label="Upload PDF or TXT Document"),
# gr.inputs.Textbox(label="Enter your prompt")
# ],
# outputs="text",
inputs=[file_input, prompt_input],
outputs=output_text,
title="Medical Report Analyzer",
description="Upload a medical report (PDF or TXT) and enter a prompt to analyze the report."
)
if __name__ == "__main__":
iface.launch(share=True)
|