Spaces:
Sleeping
Sleeping
| import logging | |
| from pathlib import Path | |
| from PyPDF2 import PdfReader | |
| from sentence_transformers import SentenceTransformer | |
| import faiss | |
| import json | |
| import concurrent.futures | |
| import numpy as np | |
| from typing import List, Tuple | |
| from datetime import datetime | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(levelname)s - %(message)s', | |
| handlers=[ | |
| logging.FileHandler('script_analysis.log'), | |
| logging.StreamHandler() | |
| ] | |
| ) | |
| class EmbeddingManager: | |
| def __init__(self, output_dir: Path, max_workers: int = 6): | |
| self.model = SentenceTransformer('all-MiniLM-L6-v2') | |
| self.chunk_size = 512 | |
| self.output_dir = output_dir | |
| (self.output_dir / 'embeddings').mkdir(exist_ok=True) | |
| self.max_workers = max_workers | |
| def extract_text(self, pdf_path: str) -> str: | |
| with open(pdf_path, 'rb') as file: | |
| reader = PdfReader(file) | |
| text = '' | |
| for page in reader.pages: | |
| text += page.extract_text() + '\n' | |
| return text | |
| def create_chunks(self, text: str) -> list: | |
| words = text.split() | |
| chunks = [] | |
| current_chunk = [] | |
| current_size = 0 | |
| for word in words: | |
| current_size += len(word) + 1 | |
| if current_size > self.chunk_size: | |
| chunks.append(' '.join(current_chunk)) | |
| current_chunk = [word] | |
| current_size = len(word) | |
| else: | |
| current_chunk.append(word) | |
| if current_chunk: | |
| chunks.append(' '.join(current_chunk)) | |
| return chunks | |
| def process_chunk_batch(self, chunks: List[str]) -> np.ndarray: | |
| """Process a batch of chunks and return their embeddings.""" | |
| try: | |
| return self.model.encode(chunks) | |
| except Exception as e: | |
| logging.error(f"Error encoding chunk batch: {e}") | |
| raise | |
| def save_embeddings(self, chunks: list,file_name:str): | |
| try: | |
| embeddings_dir = self.output_dir / 'embeddings' / file_name | |
| embeddings_dir.mkdir(parents=True, exist_ok=True) | |
| # Split chunks into batches for parallel processing | |
| batch_size = len(chunks) // self.max_workers | |
| batches = [chunks[i:i + batch_size] for i in range(0, len(chunks), batch_size)] | |
| # Process batches in parallel | |
| embeddings_list = [] | |
| with concurrent.futures.ProcessPoolExecutor(max_workers=self.max_workers) as executor: | |
| future_to_batch = {executor.submit(self.process_chunk_batch, batch): batch | |
| for batch in batches} | |
| for future in concurrent.futures.as_completed(future_to_batch): | |
| batch_embeddings = future.result() | |
| embeddings_list.append(batch_embeddings) | |
| # Combine all embeddings | |
| embeddings = np.vstack(embeddings_list) | |
| # Create and save FAISS index | |
| index = faiss.IndexFlatL2(embeddings.shape[1]) | |
| index.add(embeddings.astype('float32')) | |
| faiss.write_index(index, str(embeddings_dir / 'script.index')) | |
| # Save metadata and chunks | |
| metadata = { | |
| 'file_name': file_name, | |
| 'timestamp': datetime.now().isoformat(), | |
| 'num_chunks': len(chunks), | |
| 'chunk_size': self.chunk_size, | |
| 'model': 'all-MiniLM-L6-v2', | |
| 'embedding_dimension': embeddings.shape[1], | |
| 'num_workers': self.max_workers | |
| } | |
| with open(embeddings_dir / 'metadata.json', 'w') as f: | |
| json.dump(metadata, f, indent=4) | |
| with open(embeddings_dir / 'chunks.json', 'w') as f: | |
| json.dump(chunks, f, indent=4) | |
| logging.info(f"Saved embeddings and metadata to {embeddings_dir}") | |
| return embeddings_dir | |
| except Exception as e: | |
| logging.error(f"Error saving embeddings: {e}") | |
| raise | |
| def process_script(self, data: str,filename): | |
| chunks = self.create_chunks(data) | |
| logging.info(f"Created {len(chunks)} chunks from script") | |
| embeddings_dir = self.save_embeddings(chunks,file_name=filename) | |
| return chunks, embeddings_dir |