Spaces:
Running
Running
File size: 4,416 Bytes
0712e1f |
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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 |
import gradio as gr
import torch
from pypdf import PdfReader
from PIL import Image
import io
from transformers import (
TrOCRProcessor,
VisionEncoderDecoderModel,
AutoTokenizer,
AutoModelForCausalLM
)
# ============================================================
# Device
# ============================================================
device = "cuda" if torch.cuda.is_available() else "cpu"
# ============================================================
# Load Models (cached by HF Spaces)
# ============================================================
ocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
ocr_model = VisionEncoderDecoderModel.from_pretrained(
"microsoft/trocr-base-printed"
).to(device)
tokenizer = AutoTokenizer.from_pretrained(
"Qwen/Qwen2.5-1.5B-Instruct",
trust_remote_code=True
)
qwen_model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-1.5B-Instruct",
device_map="auto",
torch_dtype=torch.float16 if device == "cuda" else torch.float32,
trust_remote_code=True
)
# ============================================================
# Helpers
# ============================================================
def is_scanned_pdf(reader):
for page in reader.pages:
if page.extract_text():
return False
return True
def extract_text_from_pdf(file):
reader = PdfReader(file)
scanned = is_scanned_pdf(reader)
extracted_text = []
if not scanned:
# Digital PDF
for page in reader.pages:
text = page.extract_text()
if text:
extracted_text.append(text)
else:
# OCR only embedded images (HF-safe)
for page in reader.pages:
if "/XObject" in page["/Resources"]:
xobjects = page["/Resources"]["/XObject"].get_object()
for obj in xobjects:
xobj = xobjects[obj]
if xobj["/Subtype"] == "/Image":
image = Image.open(io.BytesIO(xobj.get_data())).convert("RGB")
pixel_values = ocr_processor(
images=image,
return_tensors="pt"
).pixel_values.to(device)
with torch.no_grad():
ids = ocr_model.generate(pixel_values)
text = ocr_processor.batch_decode(
ids,
skip_special_tokens=True
)[0]
extracted_text.append(text)
return "\n\n".join(extracted_text)
def evaluate_text(text):
prompt = f"""
You are a strict academic evaluator.
Evaluate the following document and assign marks out of 10.
Criteria:
- Clarity
- Structure
- Technical depth
- Language quality
- Completeness
DOCUMENT:
---------
{text[:6000]}
---------
Respond strictly in this format:
Score: X/10
Justification:
Strengths:
Weaknesses:
"""
inputs = tokenizer(prompt, return_tensors="pt").to(device)
with torch.no_grad():
output = qwen_model.generate(
**inputs,
max_new_tokens=400,
do_sample=False
)
return tokenizer.decode(output[0], skip_special_tokens=True)
# ============================================================
# Gradio Function
# ============================================================
def process_pdf(pdf_file):
extracted_text = extract_text_from_pdf(pdf_file)
evaluation = evaluate_text(extracted_text)
return extracted_text, evaluation
# ============================================================
# Gradio UI
# ============================================================
with gr.Blocks(title="PDF Evaluator (OCR + Qwen)") as demo:
gr.Markdown("""
# 📄 PDF Evaluator
Upload a PDF to:
- Extract text (OCR if needed)
- Evaluate content using Qwen
- Get marks out of 10
""")
pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
extract_btn = gr.Button("Extract & Evaluate")
extracted_output = gr.Textbox(
label="Extracted Text",
lines=20
)
evaluation_output = gr.Textbox(
label="Evaluation",
lines=10
)
extract_btn.click(
process_pdf,
inputs=pdf_input,
outputs=[extracted_output, evaluation_output]
)
demo.launch()
|