Spaces:
Sleeping
Sleeping
File size: 5,318 Bytes
f473a91 | 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 | """
Reference extraction from PDF and text documents using Google Gemini AI.
"""
from google import genai
from google.genai import types
import pathlib
import json
import logging
from typing import Dict, List, Any
logger = logging.getLogger(__name__)
class ReferenceExtractor:
"""
Extracts academic references from PDF files and text using AI.
"""
# Shared extraction prompt template
EXTRACTION_PROMPT_TEMPLATE = """
Identify the title, first author, year, and journal of each reference.
The references may be in various formats, including APA, MLA, and Chicago.
The references may also include DOIs, URLs, and other identifiers.
IMPORTANT: For each reference, analyze its content to determine the optimal API order:
- If the reference mentions "arXiv", "preprint", or contains "arxiv.org", put 'arxiv' first in the API order.
- If it mentions a specific journal, contains a DOI, or is clearly a published journal article, put 'crossref' first.
- For newer publications, open access articles, or general academic works, put 'openalex' first.
Format your response as JSON with these fields:
- title: The title of the academic work in the reference
- first_author: The first author's name
- year: Publication year as string
- journal: Journal or publication venue name (optional)
- reference: The full reference text without the number of the reference
- check_api_order: Array of API names to check in order ['crossref', 'arxiv', 'openalex'] based on the reference type
Output only json
"""
def __init__(self, client: genai.Client, model: str = "gemini-2.5-flash-lite"):
"""
Initialize the reference extractor.
Args:
client: Google GenAI client instance
model: Model name to use for extraction
"""
self.client = client
self.model = model
def extract_from_pdf(self, file_path: str) -> List[Dict[str, Any]]:
"""
Extracts references from a PDF file using GenAI.
Args:
file_path: The path to the PDF file to extract references from.
Returns:
A list of extracted references.
Raises:
FileNotFoundError: If PDF file doesn't exist
ValueError: If extraction fails
"""
file_path_obj = pathlib.Path(file_path)
if not file_path_obj.exists():
logger.error(f"PDF file not found: {file_path}")
raise FileNotFoundError(f"PDF file not found: {file_path}")
prompt = f"From the PDF file, extract the references.\n{self.EXTRACTION_PROMPT_TEMPLATE}"
try:
logger.info(f"Extracting references from PDF: {file_path_obj.name}")
response = self.client.models.generate_content(
model=self.model,
contents=[
types.Part.from_bytes(
data=file_path_obj.read_bytes(),
mime_type='application/pdf',
),
prompt
],
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0
)
)
references = json.loads(response.candidates[0].content.parts[0].text)
logger.info(f"Extracted {len(references) if isinstance(references, list) else 1} references")
return references
except json.JSONDecodeError as e:
logger.error(f"Failed to parse GenAI response as JSON: {e}")
raise ValueError(f"Invalid JSON response from GenAI: {e}")
except Exception as e:
logger.exception(f"Failed to extract references from PDF: {e}")
raise
def extract_from_text(self, text: str) -> List[Dict[str, Any]]:
"""
Extracts references from a text string using GenAI.
Args:
text: The text to extract references from.
Returns:
A list of extracted references.
Raises:
ValueError: If text is empty or extraction fails
"""
if not text or not text.strip():
logger.error("Empty text provided for reference extraction")
raise ValueError("Text cannot be empty")
prompt = f"""
From the text, extract the references.
Text:
{text}
{self.EXTRACTION_PROMPT_TEMPLATE}
"""
try:
logger.info(f"Extracting references from text ({len(text)} characters)")
response = self.client.models.generate_content(
model=self.model,
contents=[prompt],
config=types.GenerateContentConfig(
response_mime_type="application/json",
temperature=0
)
)
references = json.loads(response.candidates[0].content.parts[0].text)
logger.info(f"Extracted {len(references) if isinstance(references, list) else 1} references")
return references
except json.JSONDecodeError as e:
logger.error(f"Failed to parse GenAI response as JSON: {e}")
raise ValueError(f"Invalid JSON response from GenAI: {e}")
except Exception as e:
logger.exception(f"Failed to extract references from text: {e}")
raise
|