Spaces:
Running
Running
| import models | |
| #import constants | |
| #from langchain_experimental.text_splitter import SemanticChunker | |
| from langchain_qdrant import QdrantVectorStore, Qdrant | |
| from langchain_community.document_loaders import PyPDFLoader, UnstructuredURLLoader | |
| from qdrant_client.http.models import VectorParams | |
| import pymupdf | |
| import requests | |
| from transformers import AutoTokenizer | |
| def extract_links_from_pdf(pdf_path): | |
| links = [] | |
| doc = pymupdf.open(pdf_path) | |
| for page in doc: | |
| for link in page.get_links(): | |
| if link['uri']: | |
| links.append(link['uri']) | |
| return links | |
| def load_documents_from_url(url): | |
| try: | |
| # Check if it's a PDF | |
| if url.endswith(".pdf"): | |
| try: | |
| loader = PyPDFLoader(url) | |
| return loader.load() | |
| except Exception as e: | |
| print(f"Error loading PDF from {url}: {e}") | |
| return None | |
| # Fetch the content and check for video pages | |
| try: | |
| response = requests.head(url, timeout=10) # Timeout for fetching headers | |
| content_type = response.headers.get('Content-Type', '') | |
| except Exception as e: | |
| print(f"Error fetching headers from {url}: {e}") | |
| return None | |
| # Ignore video content (flagged for now) | |
| if 'video' in content_type: | |
| return None | |
| if 'youtube' in url: | |
| return None | |
| # Otherwise, treat it as an HTML page | |
| try: | |
| loader = UnstructuredURLLoader([url]) | |
| return loader.load() | |
| except Exception as e: | |
| print(f"Error loading HTML from {url}: {e}") | |
| return None | |
| except Exception as e: | |
| print(f"General error loading from {url}: {e}") | |
| return None | |
| #gather kai's docs | |
| filepaths = ["./test_docs/Employee Statistics FINAL.pdf","./test_docs/Employer Statistics FINAL.pdf","./test_docs/Articles To Share.pdf"] | |
| all_links = [] | |
| for pdf_path in filepaths: | |
| all_links.extend(extract_links_from_pdf(pdf_path)) | |
| unique_links = list(set(all_links)) | |
| print(unique_links) | |
| documents = [] | |
| for link in unique_links: | |
| doc = load_documents_from_url(link) | |
| #print(f"loaded doc from {link}") | |
| if doc: | |
| documents.extend(doc) | |
| #print(len(documents)) | |
| #semantic_split_docs = models.semanticChunker.split_documents(documents) | |
| semantic_tuned_split_docs = models.semanticChunker_tuned.split_documents(documents) | |
| #RCTS_split_docs = models.RCTS.split_documents(documents) | |
| #print(len(semantic_split_docs)) | |
| print(len(semantic_tuned_split_docs)) | |
| #tokenizer = models.tuned_embeddings.client.tokenizer | |
| # | |
| #token_sizes = [len(tokenizer.encode(chunk)) for chunk in semantic_tuned_split_docs] | |
| # Display the token sizes | |
| #for idx, size in enumerate(token_sizes): | |
| # print(f"Chunk {idx + 1}: {size} tokens") | |
| # | |
| #exit() | |
| #add them to the existing qdrant client | |
| collection_name = "docs_from_ripped_urls_semantic_tuned" | |
| collections = models.qdrant_client.get_collections() | |
| collection_names = [collection.name for collection in collections.collections] | |
| # If the collection does not exist, create it | |
| if collection_name not in collection_names: | |
| models.qdrant_client.create_collection( | |
| collection_name=collection_name, | |
| vectors_config=VectorParams(size=1024, distance="Cosine") | |
| ) | |
| qdrant_vector_store = QdrantVectorStore( | |
| client=models.qdrant_client, | |
| collection_name=collection_name, | |
| embedding=models.tuned_embeddings | |
| ) | |
| qdrant_vector_store.add_documents(semantic_tuned_split_docs) | |
| collection_info = models.qdrant_client.get_collection(collection_name) | |
| print(f"Number of points in collection: {collection_info.points_count}") |