import os
import os.path as osp
from typing import List
from tqdm import tqdm
from langchain.docstore.document import Document
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import NLTKTextSplitter
from langchain.vectorstores.faiss import FAISS
import pandas as pd
import nltk
nltk.download('punkt')

PROCESSED_CSV_DIRECTORY = "processed" # Directory to save processed CSV file

def create_docs() -> List[Document]:
    docs = []
    df = pd.read_csv(osp.join(PROCESSED_CSV_DIRECTORY, 'scraped.csv'))
    for index, row in df.iterrows():
        doc = Document(page_content=row['text'], metadata={"source": row['url']})
        docs.append(doc)
    return docs


docs = create_docs()

doc_chunks = []
seen_chunks = set()
total_websites = set()
total_words = 0
splitter = NLTKTextSplitter(chunk_size=1024)
for source in tqdm(docs):
    for chunk in splitter.split_text(source.page_content):
        if chunk not in seen_chunks:
            doc_chunks.append(
                Document(page_content=chunk, metadata=source.metadata))
            total_words += len(chunk.split())
            total_websites.add(source.metadata['source'])
        seen_chunks.add(chunk)

print(f'Total websites: {len(total_websites)}')
print(f'Total chunks: {len(doc_chunks)}')
print(f'Total words: {total_words}')
print(f'Avg words per chunk: {int(total_words / len(doc_chunks))}')
print(f'Estimated embedding cost: ${total_words / 0.75 / 1000 * 0.0004:.2f}')

search_index = FAISS.from_documents(doc_chunks, OpenAIEmbeddings(model='text-embedding-ada-002'))

# persistent search index
search_index.save_local("search_index")

langchain.text_splitter.NLTKTextSplitter langchain.docstore.document.Document langchain.embeddings.openai.OpenAIEmbeddings
[((316, 338), 'nltk.download', 'nltk.download', (['"""punkt"""'], {}), "('punkt')\n", (329, 338), False, 'import nltk\n'), ((810, 843), 'langchain.text_splitter.NLTKTextSplitter', 'NLTKTextSplitter', ([], {'chunk_size': '(1024)'}), '(chunk_size=1024)\n', (826, 843), False, 'from langchain.text_splitter import NLTKTextSplitter\n'), ((858, 868), 'tqdm.tqdm', 'tqdm', (['docs'], {}), '(docs)\n', (862, 868), False, 'from tqdm import tqdm\n'), ((1527, 1575), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {'model': '"""text-embedding-ada-002"""'}), "(model='text-embedding-ada-002')\n", (1543, 1575), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((490, 538), 'os.path.join', 'osp.join', (['PROCESSED_CSV_DIRECTORY', '"""scraped.csv"""'], {}), "(PROCESSED_CSV_DIRECTORY, 'scraped.csv')\n", (498, 538), True, 'import os.path as osp\n'), ((591, 658), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': "row['text']", 'metadata': "{'source': row['url']}"}), "(page_content=row['text'], metadata={'source': row['url']})\n", (599, 658), False, 'from langchain.docstore.document import Document\n'), ((1013, 1067), 'langchain.docstore.document.Document', 'Document', ([], {'page_content': 'chunk', 'metadata': 'source.metadata'}), '(page_content=chunk, metadata=source.metadata)\n', (1021, 1067), False, 'from langchain.docstore.document import Document\n')]