Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| #!/usr/bin/env python3 | |
| # scripts/create_eu_ai_act_vectorstore.py | |
| """ | |
| Script to create and save a vectorstore from the EU AI Act PDF. | |
| This creates a FAISS vectorstore that can be loaded quickly in the main app. | |
| """ | |
| import os | |
| import sys | |
| import logging | |
| from pathlib import Path | |
| import pickle | |
| from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeout | |
| from typing import Optional | |
| import dotenv | |
| # Import config | |
| from config import config | |
| # Pipeline status tracking (no-op when running standalone) | |
| _PIPELINE_SCRIPT = os.environ.get("PIPELINE_CURRENT_SCRIPT") | |
| _pipeline = None | |
| if _PIPELINE_SCRIPT: | |
| try: | |
| from pipeline_status import PipelineStatus | |
| _pipeline = PipelineStatus() | |
| except Exception: | |
| pass | |
| def _update_pipeline_progress(current, total, unit="vectorstores", message=""): | |
| if not _pipeline: | |
| return | |
| try: | |
| _pipeline.update_progress(_PIPELINE_SCRIPT, current, total, unit, message) | |
| except Exception: | |
| pass | |
| # PDF processing | |
| import PyPDF2 | |
| # LangChain components | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langchain_openai import OpenAIEmbeddings | |
| from langchain_community.vectorstores import FAISS | |
| from langchain_core.documents import Document | |
| # Load environment variables | |
| dotenv.load_dotenv() | |
| # Create logs directory if it doesn't exist | |
| os.makedirs("data_updating_scripts/logs", exist_ok=True) | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s [%(levelname)s] %(message)s", | |
| handlers=[logging.StreamHandler(), logging.FileHandler("data_updating_scripts/logs/eu_vectorstore.log")], | |
| ) | |
| logger = logging.getLogger(__name__) | |
| def extract_text_from_pdf(pdf_path: str) -> str: | |
| """Extract text from PDF file.""" | |
| try: | |
| with open(pdf_path, 'rb') as file: | |
| pdf_reader = PyPDF2.PdfReader(file) | |
| text = "" | |
| logger.info(f"Processing {len(pdf_reader.pages)} pages from {pdf_path}") | |
| for page_num, page in enumerate(pdf_reader.pages): | |
| try: | |
| page_text = page.extract_text() | |
| text += f"\n\n--- Page {page_num + 1} ---\n\n{page_text}" | |
| except Exception as e: | |
| logger.warning(f"Error extracting text from page {page_num + 1}: {e}") | |
| continue | |
| logger.info(f"Extracted {len(text)} characters from PDF") | |
| return text | |
| except Exception as e: | |
| logger.error(f"Error reading PDF {pdf_path}: {e}") | |
| raise e | |
| def create_eu_ai_act_documents(text_content: str) -> list: | |
| """Convert EU AI Act text to Document objects with metadata.""" | |
| try: | |
| # Initialize text splitter with appropriate settings for legal documents | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=1500, # Larger chunks for legal text | |
| chunk_overlap=200, # More overlap for context preservation | |
| length_function=len, | |
| separators=["\n\n", "\n", ". ", " ", ""] | |
| ) | |
| # Create initial document | |
| doc = Document( | |
| page_content=text_content, | |
| metadata={ | |
| 'source': 'EU AI Act', | |
| 'document_type': 'regulation', | |
| 'jurisdiction': 'European Union', | |
| 'title': 'Regulation (EU) 2024/1689 on Artificial Intelligence (AI Act)' | |
| } | |
| ) | |
| # Split into chunks | |
| splits = text_splitter.split_documents([doc]) | |
| # Add chunk-specific metadata | |
| for i, split in enumerate(splits): | |
| split.metadata.update({ | |
| 'chunk_id': i, | |
| 'total_chunks': len(splits) | |
| }) | |
| logger.info(f"Created {len(splits)} document chunks") | |
| return splits | |
| except Exception as e: | |
| logger.error(f"Error creating documents: {e}") | |
| raise e | |
| def create_and_save_eu_vectorstore( | |
| pdf_path: str = "data_updating_scripts/eu-ai-act.pdf", | |
| vectorstore_path: str = "data/eu_ai_act_vectorstore", | |
| openai_api_key: Optional[str] = None | |
| ) -> bool: | |
| """ | |
| Create FAISS vectorstore from EU AI Act PDF and save it locally. | |
| Args: | |
| pdf_path: Path to the EU AI Act PDF file | |
| vectorstore_path: Directory to save the vectorstore | |
| openai_api_key: OpenAI API key (if not provided, uses environment variable) | |
| Returns: | |
| bool: True if successful, False otherwise | |
| """ | |
| try: | |
| # Check if PDF exists | |
| if not Path(pdf_path).exists(): | |
| logger.error(f"PDF file not found: {pdf_path}") | |
| return False | |
| # Get API key | |
| api_key = openai_api_key or config.OPENAI_API_KEY | |
| if not api_key: | |
| logger.error("OpenAI API key not found") | |
| return False | |
| logger.info("Starting EU AI Act vectorstore creation...") | |
| # Extract text from PDF | |
| logger.info("Extracting text from PDF...") | |
| text_content = extract_text_from_pdf(pdf_path) | |
| if not text_content or len(text_content) < 1000: | |
| logger.error("Insufficient text extracted from PDF") | |
| return False | |
| # Create documents | |
| logger.info("Creating document chunks...") | |
| documents = create_eu_ai_act_documents(text_content) | |
| if not documents: | |
| logger.error("No documents created") | |
| return False | |
| # Initialize embeddings | |
| logger.info("Initializing embeddings...") | |
| embeddings = OpenAIEmbeddings( | |
| api_key=api_key, | |
| model="text-embedding-3-small", | |
| request_timeout=120, | |
| ) | |
| # Create vectorstore with thread-based timeout | |
| # (signal.alarm doesn't work — openai SDK resets SIGALRM) | |
| logger.info("Creating FAISS vectorstore...") | |
| with ThreadPoolExecutor(max_workers=1) as executor: | |
| future = executor.submit(FAISS.from_documents, documents, embeddings) | |
| try: | |
| vectorstore = future.result(timeout=300) | |
| except FuturesTimeout: | |
| raise TimeoutError("FAISS vectorstore creation exceeded 300s hard timeout") | |
| # Create directory if it doesn't exist | |
| Path(vectorstore_path).mkdir(exist_ok=True) | |
| # Save vectorstore | |
| logger.info(f"Saving vectorstore to {vectorstore_path}...") | |
| vectorstore.save_local(vectorstore_path) | |
| # Save metadata | |
| metadata = { | |
| 'pdf_path': pdf_path, | |
| 'total_chunks': len(documents), | |
| 'text_length': len(text_content), | |
| 'embedding_model': 'text-embedding-3-small', | |
| 'chunk_size': 1500, | |
| 'chunk_overlap': 200 | |
| } | |
| metadata_path = Path(vectorstore_path) / "metadata.pickle" | |
| with open(metadata_path, 'wb') as f: | |
| pickle.dump(metadata, f) | |
| logger.info(f"✅ EU AI Act vectorstore created successfully!") | |
| logger.info(f" - Total chunks: {len(documents)}") | |
| logger.info(f" - Text length: {len(text_content):,} characters") | |
| logger.info(f" - Saved to: {vectorstore_path}") | |
| return True | |
| except Exception as e: | |
| logger.error(f"Error creating EU AI Act vectorstore: {e}") | |
| return False | |
| def load_eu_vectorstore( | |
| vectorstore_path: str = "data/eu_ai_act_vectorstore", | |
| openai_api_key: Optional[str] = None | |
| ) -> Optional[FAISS]: | |
| """ | |
| Load the EU AI Act vectorstore from disk. | |
| Args: | |
| vectorstore_path: Path to the saved vectorstore | |
| openai_api_key: OpenAI API key | |
| Returns: | |
| FAISS vectorstore or None if failed | |
| """ | |
| try: | |
| if not Path(vectorstore_path).exists(): | |
| logger.error(f"Vectorstore not found: {vectorstore_path}") | |
| return None | |
| # Get API key | |
| api_key = openai_api_key or config.OPENAI_API_KEY | |
| if not api_key: | |
| logger.error("OpenAI API key not found") | |
| return None | |
| # Initialize embeddings | |
| embeddings = OpenAIEmbeddings( | |
| api_key=api_key, | |
| model="text-embedding-3-small", | |
| request_timeout=120, | |
| ) | |
| # Load vectorstore | |
| vectorstore = FAISS.load_local( | |
| vectorstore_path, | |
| embeddings, | |
| allow_dangerous_deserialization=True # Required for loading pickled objects | |
| ) | |
| logger.info(f"✅ EU AI Act vectorstore loaded from {vectorstore_path}") | |
| return vectorstore | |
| except Exception as e: | |
| logger.error(f"Error loading EU AI Act vectorstore: {e}") | |
| return None | |
| def get_vectorstore_info(vectorstore_path: str = "data/eu_ai_act_vectorstore") -> dict: | |
| """Get information about the saved vectorstore.""" | |
| try: | |
| metadata_path = Path(vectorstore_path) / "metadata.pickle" | |
| if metadata_path.exists(): | |
| with open(metadata_path, 'rb') as f: | |
| metadata = pickle.load(f) | |
| return metadata | |
| else: | |
| return {"error": "Metadata not found"} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| if __name__ == "__main__": | |
| # Create the vectorstore | |
| _update_pipeline_progress(0, 1, "vectorstores", "Building EU AI Act vectorstore...") | |
| success = create_and_save_eu_vectorstore() | |
| if success: | |
| _update_pipeline_progress(1, 1, "vectorstores", "EU AI Act vectorstore created") | |
| if success: | |
| # Display info | |
| info = get_vectorstore_info() | |
| print("\n" + "="*50) | |
| print("EU AI Act Vectorstore Information:") | |
| print("="*50) | |
| for key, value in info.items(): | |
| if key != 'error': | |
| print(f"{key}: {value}") | |
| print("="*50) | |
| sys.exit(0) # Force exit — background threads from embeddings can hang | |
| else: | |
| print("❌ Failed to create EU AI Act vectorstore") | |
| sys.exit(1) |