Spaces:
Sleeping
Sleeping
File size: 4,600 Bytes
eb4b18c | 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 | 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 |