Spaces:
Running
Running
| import gradio as gr | |
| import logging | |
| import requests | |
| import time | |
| from bs4 import BeautifulSoup | |
| from datetime import datetime | |
| from typing import List, Optional, Tuple | |
| from urllib.parse import urljoin, urlparse | |
| import random | |
| import nltk | |
| from nltk.tokenize import sent_tokenize | |
| import PyPDF2 | |
| import io | |
| from joblib import dump, load | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import numpy as np | |
| from icalendar import Calendar | |
| from requests.adapters import HTTPAdapter | |
| from requests.packages.urllib3.util.retry import Retry | |
| from fake_useragent import UserAgent | |
| from concurrent.futures import ThreadPoolExecutor | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" | |
| ) | |
| logger = logging.getLogger(__name__) | |
| # Download NLTK data | |
| try: | |
| nltk.download('punkt', quiet=True) | |
| except Exception as e: | |
| logger.warning(f"Failed to download NLTK data: {e}") | |
| class Config: | |
| MODEL_NAME = "microsoft/DialoGPT-medium" | |
| EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" | |
| MAX_TOKENS = 1000 | |
| REQUEST_TIMEOUT = 10 | |
| MAX_DEPTH = 1 | |
| SIMILARITY_THRESHOLD = 0.5 | |
| CHUNK_SIZE = 512 | |
| MAX_WORKERS = 5 | |
| INDEXED_URLS = { | |
| "https://drive.google.com/file/d/1d5kkqaQkdiA2SwJ0JFrTuKO9zauiUtFz/view?usp=sharing" | |
| } | |
| class ResourceItem: | |
| def __init__(self, url: str, content: str, resource_type: str): | |
| self.url = url | |
| self.content = content | |
| self.type = resource_type | |
| self.embedding = None | |
| self.chunks = [] | |
| self.chunk_embeddings = [] | |
| def __str__(self): | |
| return f"ResourceItem(type={self.type}, url={self.url}, content_length={len(self.content)})" | |
| def create_chunks(self, chunk_size=Config.CHUNK_SIZE): | |
| """Split content into overlapping chunks for better context preservation""" | |
| words = self.content.split() | |
| overlap = chunk_size // 4 # 25% overlap | |
| for i in range(0, len(words), chunk_size - overlap): | |
| chunk = ' '.join(words[i:i + chunk_size]) | |
| if chunk: | |
| self.chunks.append(chunk) | |
| class RobustCrawler: | |
| def __init__(self, max_retries=3, backoff_factor=0.3): | |
| self.ua = UserAgent() | |
| self.session = self._create_robust_session(max_retries, backoff_factor) | |
| def _create_robust_session(self, max_retries, backoff_factor): | |
| session = requests.Session() | |
| retry_strategy = Retry( | |
| total=max_retries, | |
| status_forcelist=[429, 500, 502, 503, 504], | |
| method_whitelist=["HEAD", "GET", "OPTIONS"], | |
| backoff_factor=backoff_factor, | |
| raise_on_status=False | |
| ) | |
| adapter = HTTPAdapter(max_retries=retry_strategy) | |
| session.mount("https://", adapter) | |
| session.mount("http://", adapter) | |
| return session | |
| def get_headers(self): | |
| return { | |
| "User-Agent": self.ua.random, | |
| "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8", | |
| "Accept-Language": "en-US,en;q=0.5", | |
| "Referer": "https://www.google.com/", | |
| "DNT": "1", | |
| "Connection": "keep-alive", | |
| "Upgrade-Insecure-Requests": "1" | |
| } | |
| def crawl_with_exponential_backoff(self, url, timeout=Config.REQUEST_TIMEOUT): | |
| try: | |
| time.sleep(random.uniform(0.5, 2.0)) | |
| response = self.session.get( | |
| url, | |
| headers=self.get_headers(), | |
| timeout=timeout | |
| ) | |
| response.raise_for_status() | |
| return response | |
| except requests.exceptions.RequestException as e: | |
| logger.error(f"Crawling error for {url}: {e}") | |
| return None | |
| class SchoolChatbot: | |
| def __init__(self): | |
| logger.info("Initializing SchoolChatbot...") | |
| self.setup_models() | |
| self.resources = [] | |
| self.visited_urls = set() | |
| self.crawl_and_index_resources() | |
| def setup_models(self): | |
| try: | |
| logger.info("Setting up models...") | |
| self.tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME) | |
| self.model = AutoModelForCausalLM.from_pretrained(Config.MODEL_NAME) | |
| self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL) | |
| logger.info("Models setup completed successfully.") | |
| except Exception as e: | |
| logger.error(f"Failed to setup models: {e}") | |
| raise RuntimeError("Failed to initialize required models") | |
| def crawl_and_index_resources(self): | |
| logger.info("Starting to crawl and index resources...") | |
| with ThreadPoolExecutor(max_workers=Config.MAX_WORKERS) as executor: | |
| futures = [executor.submit(self.crawl_url, url, 0) for url in Config.INDEXED_URLS] | |
| for future in futures: | |
| try: | |
| future.result() | |
| except Exception as e: | |
| logger.error(f"Error in crawling thread: {e}") | |
| logger.info(f"Crawling completed. Indexed {len(self.resources)} resources.") | |
| def crawl_url(self, url, depth): | |
| if depth > Config.MAX_DEPTH or url in self.visited_urls: | |
| return | |
| self.visited_urls.add(url) | |
| crawler = RobustCrawler() | |
| response = crawler.crawl_with_exponential_backoff(url) | |
| if not response: | |
| logger.error(f"Failed to retrieve content from {url}. Please check the URL and permissions.") | |
| return | |
| content_type = response.headers.get("Content-Type", "").lower() | |
| try: | |
| if "text/calendar" in content_type or url.endswith(".ics"): | |
| self.extract_ics_content(url, response.text) | |
| elif "text/html" in content_type: | |
| self.extract_html_content(url, response) | |
| elif "application/pdf" in content_type: | |
| self.extract_pdf_content(url, response.content) | |
| else: | |
| logger.warning(f"Unknown content type for {url}: {content_type}") | |
| self.store_resource(url, response.text, 'unknown') | |
| except Exception as e: | |
| logger.error(f"Error processing {url}: {e}") | |
| def extract_ics_content(self, url, ics_text): | |
| try: | |
| cal = Calendar.from_ical(ics_text) | |
| events = [] | |
| for component in cal.walk(): | |
| if component.name == "VEVENT": | |
| event = self._format_calendar_event(component) | |
| if event: | |
| events.append(event) | |
| if events: | |
| self.store_resource(url, "\n".join(events), 'calendar') | |
| except Exception as e: | |
| logger.error(f"Error parsing ICS from {url}: {e}") | |
| def _format_calendar_event(self, event): | |
| try: | |
| summary = event.get("SUMMARY", "No Summary") | |
| start = event.get("DTSTART", "").dt | |
| end = event.get("DTEND", "").dt | |
| description = event.get("DESCRIPTION", "") | |
| location = event.get("LOCATION", "") | |
| event_details = [f"Event: {summary}"] | |
| if start: | |
| event_details.append(f"Start: {start}") | |
| if end: | |
| event_details.append(f"End: {end}") | |
| if location: | |
| event_details.append(f"Location: {location}") | |
| if description: | |
| event_details.append(f"Description: {description}") | |
| return " | ".join(event_details) | |
| except Exception: | |
| return None | |
| def extract_html_content(self, url, response): | |
| try: | |
| soup = BeautifulSoup(response.content, 'html.parser') | |
| # Remove unwanted elements | |
| for element in soup.find_all(['script', 'style', 'nav', 'footer']): | |
| element.decompose() | |
| content_sections = [] | |
| # Extract main content | |
| main_content = soup.find(['main', 'article', 'div'], class_=['content', 'main-content']) | |
| if main_content: | |
| content_sections.append(main_content.get_text(strip=True, separator=' ')) | |
| # Extract headings and their associated content | |
| for heading in soup.find_all(['h1', 'h2', 'h3']): | |
| section = [heading.get_text(strip=True)] | |
| next_elem = heading.find_next_sibling() | |
| while next_elem and next_elem.name in ['p', 'ul', 'ol', 'div']: | |
| section.append(next_elem.get_text(strip=True)) | |
| next_elem = next_elem.find_next_sibling() | |
| content_sections.append(' '.join(section)) | |
| if content_sections: | |
| self.store_resource(url, ' '.join(content_sections), 'webpage') | |
| # Process links if within depth limit | |
| if len(self.visited_urls) < Config.MAX_DEPTH: | |
| self._process_links(soup, url) | |
| except Exception as e: | |
| logger.error(f"Error extracting HTML content from {url}: {e}") | |
| def _process_links(self, soup, base_url): | |
| try: | |
| for link in soup.find_all('a', href=True): | |
| full_url = urljoin(base_url, link['href']) | |
| if self.is_valid_url(full_url) and full_url not in self.visited_urls: | |
| time.sleep(random.uniform(0.5, 2.0)) | |
| self.crawl_url(full_url, len(self.visited_urls)) | |
| except Exception as e: | |
| logger.error(f"Error processing links from {base_url}: {e}") | |
| def extract_pdf_content(self, url, pdf_content): | |
| try: | |
| pdf_file = io.BytesIO(pdf_content) | |
| pdf_reader = PyPDF2.PdfReader(pdf_file) | |
| text_content = [] | |
| for page in pdf_reader.pages: | |
| try: | |
| text_content.append(page.extract_text()) | |
| except Exception as e: | |
| logger.error(f"Error extracting text from PDF page: {e}") | |
| continue | |
| if text_content: | |
| self.store_resource(url, ' '.join(text_content), 'pdf') | |
| except Exception as e: | |
| logger.error(f"Error extracting PDF content from {url}: {e}") | |
| def store_resource(self, url, text_data, resource_type): | |
| try: | |
| # Create resource item and split into chunks | |
| item = ResourceItem(url, text_data, resource_type) | |
| item.create_chunks() | |
| # Generate embeddings for chunks | |
| item.chunk_embeddings = [ | |
| self.embedding_model.encode(chunk) | |
| for chunk in item.chunks | |
| ] | |
| # Calculate average embedding | |
| if item.chunk_embeddings: | |
| item.embedding = np.mean(item.chunk_embeddings, axis=0) | |
| self.resources.append(item) | |
| logger.debug(f"Stored resource: {url} (type={resource_type})") | |
| except Exception as e: | |
| logger.error(f"Error storing resource {url}: {e}") | |
| def is_valid_url(self, url): | |
| try: | |
| parsed = urlparse(url) | |
| return bool(parsed.scheme) and bool(parsed.netloc) | |
| except Exception: | |
| return False | |
| def find_best_matching_chunks(self, query, n_chunks=3): | |
| if not self.resources: | |
| return [] | |
| try: | |
| query_embedding = self.embedding_model.encode(query) | |
| all_chunks = [] | |
| for resource in self.resources: | |
| for chunk, embedding in zip(resource.chunks, resource.chunk_embeddings): | |
| score = cosine_similarity([query_embedding], [embedding])[0][0] | |
| if score > Config.SIMILARITY_THRESHOLD: | |
| all_chunks.append((chunk, score, resource.url)) | |
| # Sort by similarity score and get top n chunks | |
| all_chunks.sort(key=lambda x: x[1], reverse=True) | |
| return all_chunks[:n_chunks] | |
| except Exception as e: | |
| logger.error(f"Error finding matching chunks: {e}") | |
| return [] | |
| def generate_response(self, user_input): | |
| try: | |
| # Find best matching chunks | |
| best_chunks = self.find_best_matching_chunks(user_input) | |
| if not best_chunks: | |
| return "I apologize, but I couldn't find any relevant information in my knowledge base. Could you please rephrase your question or ask about something else?" | |
| # Prepare context from best matching chunks | |
| context = "\n".join([chunk[0] for chunk in best_chunks]) | |
| # Prepare conversation history | |
| conversation = f"Context: {context}\nUser: {user_input}\nAssistant:" | |
| # Generate response | |
| input_ids = self.tokenizer.encode(conversation, return_tensors='pt') | |
| response_ids = self.model.generate( | |
| input_ids, | |
| max_length=Config.MAX_TOKENS, | |
| pad_token_id=self.tokenizer.eos_token_id, | |
| temperature=0.7, | |
| top_p=0.9, | |
| do_sample=True | |
| ) | |
| response = self.tokenizer.decode( | |
| response_ids[:, input_ids.shape[-1]:][0], | |
| skip_special_tokens=True | |
| ) | |
| # Format response with source | |
| source_urls = list(set(chunk[2] for chunk in best_chunks)) | |
| sources = "\n\nSources:\n" + "\n".join(source_urls) | |
| return response + sources | |
| except Exception as e: | |
| logger.error(f"Error generating response: {e}") | |
| return "I apologize, but I encountered an error while processing your question. Please try again." | |
| def create_gradio_interface(chatbot): | |
| def respond(user_input): | |
| return chatbot.generate_response(user_input) | |
| interface = gr.Interface( | |
| fn=respond, | |
| inputs=gr.Textbox( | |
| label="Ask a Question", | |
| placeholder="Type your question here...", | |
| lines=2 | |
| ), | |
| outputs=gr.Textbox( | |
| label="Answer", | |
| placeholder="Response will appear here...", | |
| lines=5 | |
| ), | |
| title="School Information Chatbot", | |
| description="Ask about school events, policies, or other information. The chatbot will provide answers based on available school documents and resources.", | |
| examples=[ | |
| ["What events are happening this week?"], | |
| ["When is the next board meeting?"], | |
| ["What is the school's attendance policy?"] | |
| ], | |
| theme=gr.themes.Soft(), | |
| flagging_mode="never" | |
| ) | |
| return interface | |
| if __name__ == "__main__": | |
| try: | |
| chatbot = SchoolChatbot() | |
| interface = create_gradio_interface(chatbot) | |
| interface.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| share=False, | |
| debug=True | |
| ) | |
| except Exception as e: | |
| logger.error(f"Failed to start application: {e}") | |