Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import pandas as pd | |
| import numpy as np | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| import os | |
| from datetime import datetime | |
| import requests | |
| from datasets import load_dataset | |
| from urllib.parse import quote | |
| # Initialize session state | |
| if 'search_history' not in st.session_state: | |
| st.session_state['search_history'] = [] | |
| if 'search_columns' not in st.session_state: | |
| st.session_state['search_columns'] = [] | |
| if 'initial_search_done' not in st.session_state: | |
| st.session_state['initial_search_done'] = False | |
| if 'dataset' not in st.session_state: | |
| st.session_state['dataset'] = None | |
| class DatasetSearcher: | |
| def __init__(self, dataset_id="tomg-group-umd/cinepile"): | |
| self.dataset_id = dataset_id | |
| self.text_model = SentenceTransformer('all-MiniLM-L6-v2') | |
| self.token = os.environ.get('DATASET_KEY') | |
| if not self.token: | |
| st.error("Please set the DATASET_KEY environment variable with your Hugging Face token.") | |
| st.stop() | |
| self.load_dataset() | |
| def load_dataset(self): | |
| """Load dataset using the datasets library""" | |
| try: | |
| if st.session_state['dataset'] is None: | |
| with st.spinner("Loading dataset..."): | |
| st.session_state['dataset'] = load_dataset( | |
| self.dataset_id, | |
| token=self.token, | |
| streaming=False | |
| ) | |
| self.dataset = st.session_state['dataset'] | |
| # Convert first split to DataFrame for easier processing | |
| first_split = next(iter(self.dataset.values())) | |
| self.df = pd.DataFrame(first_split) | |
| # Store column information | |
| self.columns = list(self.df.columns) | |
| self.text_columns = [col for col in self.columns | |
| if self.df[col].dtype == 'object' | |
| and not any(term in col.lower() | |
| for term in ['embed', 'vector', 'encoding'])] | |
| # Update session state columns | |
| st.session_state['search_columns'] = self.text_columns | |
| # Prepare text embeddings | |
| self.prepare_features() | |
| except Exception as e: | |
| st.error(f"Error loading dataset: {str(e)}") | |
| st.error("Please check your authentication token and internet connection.") | |
| st.stop() | |
| def prepare_features(self): | |
| """Prepare text embeddings for semantic search""" | |
| try: | |
| # Combine text columns for embedding | |
| combined_text = self.df[self.text_columns].fillna('').agg(' '.join, axis=1) | |
| # Create embeddings in batches to manage memory | |
| batch_size = 32 | |
| all_embeddings = [] | |
| with st.spinner("Preparing search features..."): | |
| for i in range(0, len(combined_text), batch_size): | |
| batch = combined_text[i:i+batch_size].tolist() | |
| embeddings = self.text_model.encode(batch) | |
| all_embeddings.append(embeddings) | |
| self.text_embeddings = np.vstack(all_embeddings) | |
| except Exception as e: | |
| st.warning(f"Error preparing features: {str(e)}") | |
| self.text_embeddings = np.random.randn(len(self.df), 384) | |
| def search(self, query, column=None, top_k=20): | |
| """Search the dataset using semantic and keyword matching""" | |
| if self.df.empty: | |
| return [] | |
| # Get semantic similarity scores | |
| query_embedding = self.text_model.encode([query])[0] | |
| similarities = cosine_similarity([query_embedding], self.text_embeddings)[0] | |
| # Get keyword match scores | |
| search_columns = [column] if column and column != "All Fields" else self.text_columns | |
| keyword_scores = np.zeros(len(self.df)) | |
| for col in search_columns: | |
| if col in self.df.columns: | |
| matches = self.df[col].fillna('').str.lower().str.count(query.lower()) | |
| keyword_scores += matches | |
| # Combine scores | |
| combined_scores = 0.5 * similarities + 0.5 * (keyword_scores / max(1, keyword_scores.max())) | |
| # Get top results | |
| top_k = min(top_k, len(combined_scores)) | |
| top_indices = np.argsort(combined_scores)[-top_k:][::-1] | |
| # Format results | |
| results = [] | |
| for idx in top_indices: | |
| result = { | |
| 'relevance_score': float(combined_scores[idx]), | |
| 'semantic_score': float(similarities[idx]), | |
| 'keyword_score': float(keyword_scores[idx]), | |
| **self.df.iloc[idx].to_dict() | |
| } | |
| results.append(result) | |
| return results | |
| def get_dataset_info(self): | |
| """Get information about the dataset""" | |
| if not self.dataset: | |
| return {} | |
| info = { | |
| 'splits': list(self.dataset.keys()), | |
| 'total_rows': sum(split.num_rows for split in self.dataset.values()), | |
| 'columns': self.columns, | |
| 'text_columns': self.text_columns, | |
| 'sample_rows': len(self.df), | |
| 'embeddings_shape': self.text_embeddings.shape | |
| } | |
| return info | |
| def render_video_result(result): | |
| """Render a video result with enhanced display""" | |
| col1, col2 = st.columns([2, 1]) | |
| with col1: | |
| if 'title' in result: | |
| st.markdown(f"**Title:** {result['title']}") | |
| if 'description' in result: | |
| st.markdown("**Description:**") | |
| st.write(result['description']) | |
| # Show timing information if available | |
| if 'start_time' in result and 'end_time' in result: | |
| st.markdown(f"**Time Range:** {result['start_time']}s - {result['end_time']}s") | |
| # Show additional metadata | |
| for key, value in result.items(): | |
| if key not in ['title', 'description', 'start_time', 'end_time', 'duration', | |
| 'relevance_score', 'semantic_score', 'keyword_score', | |
| 'video_id', 'youtube_id']: | |
| st.markdown(f"**{key.replace('_', ' ').title()}:** {value}") | |
| with col2: | |
| # Show search scores | |
| st.markdown("**Search Scores:**") | |
| cols = st.columns(3) | |
| cols[0].metric("Overall", f"{result['relevance_score']:.2%}") | |
| cols[1].metric("Semantic", f"{result['semantic_score']:.2%}") | |
| cols[2].metric("Keyword", f"{result['keyword_score']:.0f} matches") | |
| # Display video if available | |
| if 'youtube_id' in result: | |
| st.video(f"https://youtube.com/watch?v={result['youtube_id']}&t={result.get('start_time', 0)}") | |
| def main(): | |
| st.title("🎥 Video Dataset Search") | |
| # Initialize search class | |
| searcher = DatasetSearcher() | |
| # Create tabs | |
| tab1, tab2 = st.tabs(["🔍 Search", "📊 Dataset Info"]) | |
| # ---- Tab 1: Search ---- | |
| with tab1: | |
| st.subheader("Search Videos") | |
| col1, col2 = st.columns([3, 1]) | |
| with col1: | |
| query = st.text_input("Search query:", | |
| value="" if st.session_state['initial_search_done'] else "") | |
| with col2: | |
| search_column = st.selectbox("Search in field:", | |
| ["All Fields"] + st.session_state['search_columns']) | |
| col3, col4 = st.columns(2) | |
| with col3: | |
| num_results = st.slider("Number of results:", 1, 100, 20) | |
| with col4: | |
| search_button = st.button("🔍 Search") | |
| if search_button and query: | |
| st.session_state['initial_search_done'] = True | |
| selected_column = None if search_column == "All Fields" else search_column | |
| with st.spinner("Searching..."): | |
| results = searcher.search(query, selected_column, num_results) | |
| st.session_state['search_history'].append({ | |
| 'query': query, | |
| 'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"), | |
| 'results': results[:5] | |
| }) | |
| for i, result in enumerate(results, 1): | |
| with st.expander( | |
| f"Result {i}: {result.get('title', result.get('description', 'No title'))[:100]}...", | |
| expanded=(i==1) | |
| ): | |
| render_video_result(result) | |
| # ---- Tab 2: Dataset Info ---- | |
| with tab2: | |
| st.subheader("Dataset Information") | |
| info = searcher.get_dataset_info() | |
| if info: | |
| st.write(f"### Dataset: {searcher.dataset_id}") | |
| st.write(f"- Total rows: {info['total_rows']:,}") | |
| st.write(f"- Available splits: {', '.join(info['splits'])}") | |
| st.write(f"- Number of columns: {len(info['columns'])}") | |
| st.write(f"- Searchable text columns: {', '.join(info['text_columns'])}") | |
| st.write("### Sample Data") | |
| st.dataframe(searcher.df.head()) | |
| st.write("### Column Details") | |
| for col in info['columns']: | |
| st.write(f"- **{col}**: {searcher.df[col].dtype}") | |
| # Sidebar | |
| with st.sidebar: | |
| st.subheader("⚙️ Search History") | |
| if st.button("🗑️ Clear History"): | |
| st.session_state['search_history'] = [] | |
| st.experimental_rerun() | |
| st.markdown("### Recent Searches") | |
| for entry in reversed(st.session_state['search_history'][-5:]): | |
| with st.expander(f"{entry['timestamp']}: {entry['query']}"): | |
| for i, result in enumerate(entry['results'], 1): | |
| st.write(f"{i}. {result.get('title', result.get('description', 'No title'))[:100]}...") | |
| if __name__ == "__main__": | |
| main() |