Update app.py

app.py

@@ -1,84 +1,64 @@
+import requests
+from bs4 import BeautifulSoup
+from urllib.parse import urljoin, urlparse
+import json
 import streamlit as st
-from langchain_groq import ChatGroq
-from langchain_community.document_loaders import WebBaseLoader
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains.combine_documents import create_stuff_documents_chain
-from langchain_core.prompts import ChatPromptTemplate
-from langchain.chains import create_retrieval_chain
-from langchain_community.vectorstores import FAISS
-from transformers import AutoTokenizer, AutoModel
-import torch
-import numpy as np
-import time
 
-# Load tokenizer and model for embeddings
-try:
-    tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-    model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-except ImportError as e:
-    st.error(f"ImportError: {e}. Make sure you have the transformers library installed.")
-except Exception as e:
-    st.error(f"Failed to load tokenizer or model: {e}")
+# Function to crawl a web page
+def crawl(base_url, depth):
+    visited = set()
+    queue = [(base_url, 0)]
+    results = []
+    base_netloc = urlparse(base_url).netloc
 
-class CustomHuggingFaceEmbeddings(HuggingFaceEmbeddings):
-    def __init__(self):
-        super().__init__(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    while queue:
+        current_url, current_depth = queue.pop(0)
+        if current_depth > depth:
+            continue
 
-    def embed_documents(self, texts):
-        inputs = tokenizer(texts, padding=True, truncation=True, return_tensors='pt')
-        with torch.no_grad():
-            embeddings = model(**inputs).last_hidden_state.mean(dim=1)
-        return embeddings.numpy()
+        if current_url in visited:
+            continue
+        
+        visited.add(current_url)
+        try:
+            response = requests.get(current_url)
+            soup = BeautifulSoup(response.content, 'html.parser')
+            text = soup.get_text()
 
-# Instantiate embeddings class
-embeddings = CustomHuggingFaceEmbeddings()
+            results.append({'url': current_url, 'content': text})
 
-if "vector" not in st.session_state:
-    st.session_state.loader = WebBaseLoader("https://docs.nvidia.com/cuda/")
-    st.session_state.docs = st.session_state.loader.load()
+            # Find all links on the page
+            for link in soup.find_all('a', href=True):
+                href = link['href']
+                full_url = urljoin(current_url, href)
+                # Check if the link is within the base domain
+                if urlparse(full_url).netloc == base_netloc:
+                    if full_url not in visited:
+                        queue.append((full_url, current_depth + 1))
 
-    st.session_state.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    documents = st.session_state.text_splitter.split_documents(st.session_state.docs[:50])
+        except Exception as e:
+            print(f"Failed to fetch {current_url}: {e}")
 
-    # Create FAISS index using the custom embeddings class
-    texts = [doc.page_content for doc in documents]
-    embedded_texts = embeddings.embed_documents(texts)
-    faiss_input = [(embedding, text) for embedding, text in zip(embedded_texts, texts)]
-    st.session_state.vectors = FAISS.from_embeddings(faiss_input)
+    return results
 
-st.title("ChatGroq Demo")
-groq_api_key = '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'
-llm = ChatGroq(groq_api_key=groq_api_key, model_name="mixtral-8x7b-32768")
+# Streamlit application
+st.title("Custom Web Crawler Demo")
 
-prompt = ChatPromptTemplate.from_template(
-    """
-    Answer the questions based on the provided context only.
-    Please provide the most accurate response based on the question
-    <context>
-    {context}
-    <context>
-    Questions: {input}
-    """
-)
+depth = st.slider("Depth", min_value=1, max_value=5, value=2)
+base_url = st.text_input("Enter Base URL", "https://docs.nvidia.com/cuda/")
 
-document_chain = create_stuff_documents_chain(llm, prompt)
-retriever = st.session_state.vectors.as_retriever()
-retrieval_chain = create_retrieval_chain(retriever, document_chain)
+if st.button("Crawl"):
+    with st.spinner('Crawling...'):
+        data = crawl(base_url, depth)
+        st.write(f"Found {len(data)} pages")
 
-user_prompt = st.text_input("Input your prompt here")
+        # Optionally save the results to a JSON file
+        with open('crawled_data.json', 'w') as f:
+            json.dump(data, f, indent=2)
 
-if user_prompt:
-    start = time.process_time()
-    try:
-        response = retrieval_chain.invoke({"input": user_prompt})
-        print("Response time:", time.process_time() - start)
-        st.write(response['answer'])
+        st.write(data)
 
-        # With a Streamlit expander
-        with st.expander("Document Similarity Search"):
-            for doc in response.get("context", []):
-                st.write(doc.page_content)
-                st.write("--------------------------------")
-    except Exception as e:
-        st.error(f"Error during retrieval or response generation: {e}")
+        # Display the first page's content for demo purposes
+        if data:
+            st.write("First page content:")
+            st.write(data[0]['content'])