# import gradio as gr import os import torch from transformers import AutoProcessor, MllamaForConditionalGeneration, TextIteratorStreamer from PIL import Image import spaces import tempfile import requests from PyPDF2 import PdfReader from threading import Thread from flask import Flask, request, jsonify import io import fitz # Check if we're running in a Hugging Face Space and if SPACES_ZERO_GPU is enabled # IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1" # IS_SPACE = os.environ.get("SPACE_ID", None) is not None # Determine the device (GPU if available, else CPU) device = "cuda" if torch.cuda.is_available() else "cpu" LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1" print(f"Using device: {device}") print(f"Low memory mode: {LOW_MEMORY}") app = Flask(__name__) # Get Hugging Face token from environment variables HF_TOKEN = os.environ.get('HF_TOKEN') # Load the model and processor model_name = "meta-llama/Llama-3.2-11B-Vision-Instruct" model = MllamaForConditionalGeneration.from_pretrained( model_name, use_auth_token=HF_TOKEN, torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32, device_map="auto" if device == "cuda" else None, # Use device mapping if CUDA is available ) # Move the model to the appropriate device (GPU if available) # model.to(device) processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN) def extract_image_from_pdf(pdf_url, dpi=75): """ Extract first page of PDF as image in memory Args: pdf_url (str): URL of PDF dpi (int): Image resolution Returns: PIL.Image: First page as image or None """ try: # Download PDF response = requests.get(pdf_url, timeout=30) response.raise_for_status() # Open PDF from bytes pdf_document = fitz.open(stream=response.content, filetype="pdf") # Get first page first_page = pdf_document[0] # Render page to pixmap pix = first_page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72)) # Convert to PIL Image img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) pdf_document.close() return img except Exception as e: print(f"Error extracting first page: {e}") return None def predict_image(image_url, text, file_pref): try: # Download the image from the URL # response = requests.get(image_url) # response.raise_for_status() # Raise an error for invalid responses # image = Image.open(io.BytesIO(response.content)).convert("RGB") if file_pref == 'img': response = requests.get(image_url) response.raise_for_status() # Raise an error for invalid responses image = Image.open(io.BytesIO(response.content)).convert("RGB") else: image = extract_image_from_pdf(image_url) messages = [ {"role": "user", "content": [ {"type": "image"}, # Specify that an image is provided {"type": "text", "text": text} # Add the user-provided text input ]} ] # Create the input text using the processor's chat template input_text = processor.apply_chat_template(messages, add_generation_prompt=True) # Process the inputs and move to the appropriate device inputs = processor(image, input_text, return_tensors="pt").to(device) # outputs = model.generate(**inputs, max_new_tokens=100) # # Decode the output to return the final response # response = processor.decode(outputs[0], skip_special_tokens=True) # return response streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True) generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=4096) generated_text = "" thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() buffer = "" for new_text in streamer: buffer += new_text # generated_text_without_prompt = buffer # # time.sleep(0.01) # yield buffer return buffer except Exception as e: raise ValueError(f"Error during prediction: {str(e)}") def extract_text_from_pdf(pdf_url): try: response = requests.get(pdf_url) response.raise_for_status() with tempfile.NamedTemporaryFile(delete=False) as temp_pdf: temp_pdf.write(response.content) temp_pdf_path = temp_pdf.name reader = PdfReader(temp_pdf_path) text = "" for page in reader.pages: text += page.extract_text() os.remove(temp_pdf_path) return text except Exception as e: raise ValueError(f"Error extracting text from PDF: {str(e)}") # raise HTTPException(status_code=400, detail=f"Error extracting text from PDF: {str(e)}") @spaces.GPU def predict_text(text): # pdf_text = extract_text_from_pdf('https://arinsight.co/2024_FA_AEC_1200_GR1_GR2.pdf') text_combined = text # + "\n\nExtracted Text from PDF:\n" + pdf_text # Prepare the input messages messages = [{"role": "user", "content": [{"type": "text", "text": text_combined}]}] # Create the input text using the processor's chat template input_text = processor.apply_chat_template(messages, add_generation_prompt=True) # Process the inputs and move to the appropriate device # inputs = processor(image, input_text, return_tensors="pt").to(device) inputs = processor(text=input_text, return_tensors="pt").to("cuda") # Generate a response from the model # outputs = model.generate(**inputs, max_new_tokens=1024) # # Decode the output to return the final response # response = processor.decode(outputs[0], skip_special_tokens=True, skip_prompt=True) streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True) generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=2048) generated_text = "" thread = Thread(target=model.generate, kwargs=generation_kwargs) thread.start() buffer = "" for new_text in streamer: buffer += new_text # generated_text_without_prompt = buffer # # time.sleep(0.01) # yield buffer return buffer PROMPT = ( "Extract the following information as per this format:\n" "'Course Code:'\n" "'Course Name:'\n" "'Course Description:'\n" "'Course Credits:'\n" "'Course Learning Outcomes:'\n" "'Delivery Method:'\n" "'Prerequisite(s):'\n" "'Co-requisite(s):'\n" "'Materials:'\n" "'Topical Outline:'\n" "Do not add anything else except the required information from this text." ) PROMPT_SKILLS = ( "Provide skills based on the Lightcast Open Skills Taxonomy in categories as:\n" "'Primary Skills' (the degree program or certification),\n" "'Secondary Skills', and\n" "'Tertiary Skills'." ) # PROMPT_IMAGE = ( # "You are a highly intelligent assistant designed to analyze images and extract structured information from them. " # "Your task is to analyze the given image of a student's academic record and generate a response in the exact JSON format provided below. " # "If any specific information is missing or unavailable in the image, replace the corresponding field with null. " # "Ensure the format is consistent, strictly adhering to the structure shown below.\n\n" # "Required JSON Format:\n\n" # "{\n" # ' "student": {\n' # ' "name": "string",\n' # ' "id": "string",\n' # ' "dob": "string",\n' # ' "original_start_date": "string",\n' # ' "cumulative_gpa": "string",\n' # ' "program": "string",\n' # ' "status": "string"\n' # ' },\n' # ' "courses": [\n' # ' {\n' # ' "transfer_institution": "string",\n' # ' "course_code": "string",\n' # ' "course_name": "string",\n' # ' "credits_attempted": number,\n' # ' "credits_earned": number,\n' # ' "grade": "string",\n' # ' "quality_points": number,\n' # ' "semester_code": "string",\n' # ' "semester_dates": "string"\n' # ' }\n' # " // Additional courses can be added here\n" # " ]\n" # "}\n\n" # "Instructions:\n\n" # "1. Extract the student information and course details as displayed in the image.\n" # "2. Use null for any missing or unavailable information.\n" # "3. Format the extracted data exactly as shown above. Do not deviate from this structure.\n" # "4. Use accurate field names and ensure proper nesting of data (e.g., 'student' and 'courses' sections).\n" # "5. The values for numeric fields like credits_attempted, credits_earned, and quality_points should be numbers (not strings).\n" # ) PROMPT_IMAGE_STUDENT = ( "You are a highly intelligent assistant designed to analyze images and extract structured information from them. " "Your task is to analyze the given image of a student's academic record and generate a response in the exact JSON format provided below. " "If any specific information is missing or unavailable in the image, replace the corresponding field with null. " "Ensure the format is consistent, strictly adhering to the structure shown below.\n\n" "Required JSON Format:\n\n" "{\n" ' "student": {\n' ' "name": "string",\n' ' "id": "string",\n' ' "dob": "string",\n' ' "original_start_date": "string",\n' ' "cumulative_gpa": "string",\n' ' "program": "string",\n' ' "status": "string"\n' ' }\n' "}\n\n" "Instructions:\n\n" "1. Extract the student's general information as displayed in the image.\n" "2. Use null for any missing or unavailable information.\n" "3. Format the extracted data exactly as shown above. Do not deviate from this structure.\n" "4. Ensure accurate field names and proper nesting.\n" "5. Return only the 'student' section as JSON.\n" ) PROMPT_IMAGE_COURSES = ( "You are a highly intelligent assistant designed to analyze images and extract structured information from them. " "Your task is to analyze the given image of a student's academic record and generate a response in the exact JSON format provided below. " "If any specific information is missing or unavailable in the image, replace the corresponding field with null. " "Ensure the format is consistent, strictly adhering to the structure shown below.\n\n" "Required JSON Format:\n\n" "{\n" ' "courses": [\n' ' {\n' ' "transfer_institution": "string",\n' ' "course_code": "string",\n' ' "course_name": "string",\n' ' "credits_attempted": number,\n' ' "credits_earned": number,\n' ' "grade": "string",\n' ' "quality_points": number,\n' ' "semester_code": "string",\n' ' "semester_dates": "string"\n' ' }\n' " // Additional courses can be added here\n" " ]\n" "}\n\n" "Instructions:\n\n" "1. Extract the course details as displayed in the image.\n" "2. Use null for any missing or unavailable information.\n" "3. Format the extracted data exactly as shown above. Do not deviate from this structure.\n" "4. Ensure accurate field names and proper nesting.\n" "5. Return only the 'courses' section as JSON.\n" ) @app.route("/", methods=["GET"]) def home(): return jsonify({"message": "Welcome to the PDF Extraction API. Use the /extract endpoint to extract information."}) @app.route("/favicon.ico") def favicon(): return "", 204 @app.route("/extract", methods=["POST"]) def extract_info(): data = request.json if not data: return jsonify({"error": "Please provide a PDF URL in the request body."}), 400 try: if data["url"] is not None: pdf_url = data["url"] pdf_text = extract_text_from_pdf(pdf_url) prompt = f"{PROMPT}\n\n{pdf_text}" response = predict_text(prompt) else: response = '' if data["skills"] == True: if response: prompt_skills = f"{PROMPT_SKILLS} using this information only -- {response}" response_skills = predict_text(prompt_skills) else: response_skills = '' else: response_skills = '' if data["img_url"] is not None: prompt_student = f"{PROMPT_IMAGE_STUDENT}\n" prompt_courses = f"{PROMPT_IMAGE_COURSES}\n" img_url = data["img_url"] file_pref = data["file_pref"] response_student = predict_image(img_url, prompt_student, file_pref) response_courses = predict_image(img_url, prompt_courses, file_pref) # response_image = response_student + response_courses response_image = {"student": response_student.get("student", {}), "courses": response_courses.get("courses", [])} else: response_image = '' return jsonify({"extracted_info": response + "\n" + response_skills + "\n" + response_image}) except Exception as e: return jsonify({"error": str(e)}), 500 if __name__ == "__main__": app.run(host="0.0.0.0", port=7860) # # Define the Gradio interface # interface = gr.Interface( # fn=predict_text, # inputs=[ # # gr.Image(type="pil", label="Image Input"), # Image input with label # gr.Textbox(label="Text Input") # Textbox input with label # ], # outputs=gr.Textbox(label="Generated Response"), # Output with a more descriptive label # title="Llama 3.2 11B Vision Instruct Demo", # Title of the interface # description="This demo uses Meta's Llama 3.2 11B Vision model to generate responses based on an image and text input.", # Short description # theme="compact" # Using a compact theme for a cleaner look # ) # # Launch the interface # interface.launch(debug=True)