app.py

import streamlit as st
import tensorflow as tf
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import pandas as pd
import os
import time
import sentencepiece as spm

# Set page title
st.set_page_config(page_title="Embedding Model Comparison", layout="wide")

# Function to load the SentencePiece tokenizer
@st.cache_resource
def load_tokenizer(tokenizer_path="sentencepiece.model"):
    if not os.path.exists(tokenizer_path):
        st.error(f"Tokenizer file not found: {tokenizer_path}")
        return None
    
    sp = spm.SentencePieceProcessor()
    sp.load(tokenizer_path)
    return sp

# Function to load a TFLite model
def load_model(model_path):
    if not os.path.exists(model_path):
        st.error(f"Model file not found: {model_path}")
        return None
    
    interpreter = tf.lite.Interpreter(model_path=model_path)
    interpreter.allocate_tensors()
    return interpreter

# Function to get embeddings from a TFLite model
def get_embedding(text, interpreter, tokenizer):
    if interpreter is None or tokenizer is None:
        return None, 0
    
    # Get input and output details
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()
    
    # Get the expected input shape
    input_shape = input_details[0]['shape']
    max_seq_length = input_shape[1] if len(input_shape) > 1 else 64
    
    # Properly tokenize the text using SentencePiece
    tokens = tokenizer.encode(text, out_type=int)
    
    # Handle padding/truncation
    if len(tokens) > max_seq_length:
        tokens = tokens[:max_seq_length]  # Truncate
    else:
        tokens = tokens + [0] * (max_seq_length - len(tokens))  # Pad
    
    # Prepare input tensor with proper shape
    token_ids = np.array([tokens], dtype=np.int32)
    
    # Set input tensor
    interpreter.set_tensor(input_details[0]['index'], token_ids)
    
    # Run inference
    start_time = time.time()
    interpreter.invoke()
    inference_time = time.time() - start_time
    
    # Get output tensor
    embedding = interpreter.get_tensor(output_details[0]['index'])
    
    return embedding, inference_time

# Function to load sentences from a file
def load_sentences(file_path):
    if not os.path.exists(file_path):
        return ["Hello world", "This is a test", "Embedding models are useful", 
                "TensorFlow Lite is great for mobile applications",
                "Streamlit makes it easy to create web apps",
                "Python is a popular programming language",
                "Machine learning is an exciting field",
                "Natural language processing helps computers understand human language",
                "Semantic search finds meaning, not just keywords",
                "Quantization reduces model size with minimal accuracy loss"]
    
    with open(file_path, 'r') as f:
        sentences = [line.strip() for line in f if line.strip()]
    
    return sentences

# Function to find similar sentences
def find_similar_sentences(query_embedding, sentence_embeddings, sentences):
    if query_embedding is None or len(sentence_embeddings) == 0:
        return []
    
    # Calculate similarity scores
    similarities = cosine_similarity(query_embedding, sentence_embeddings)[0]
    
    # Get indices sorted by similarity (descending)
    sorted_indices = np.argsort(similarities)[::-1]
    
    # Create result list
    results = []
    for idx in sorted_indices:
        results.append({
            "sentence": sentences[idx],
            "similarity": similarities[idx]
        })
    
    return results

# Main application
def main():
    st.title("Embedding Model Comparison")
    
    # Sidebar for configuration
    with st.sidebar:
        st.header("Configuration")
        old_model_path = st.text_input("Old Model Path", "old.tflite")
        new_model_path = st.text_input("New Model Path", "new.tflite")
        sentences_path = st.text_input("Sentences File Path", "sentences.txt")
        tokenizer_path = st.text_input("Tokenizer Path", "sentencepiece.model")
    
    # Load the tokenizer
    tokenizer = load_tokenizer(tokenizer_path)
    if tokenizer:
        st.sidebar.success("Tokenizer loaded successfully")
        st.sidebar.write(f"Vocabulary size: {tokenizer.get_piece_size()}")
    else:
        st.sidebar.error("Failed to load tokenizer")
        return
    
    # Load the models
    st.header("Models")
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("Old Model")
        old_model = load_model(old_model_path)
        if old_model:
            st.success("Old model loaded successfully")
            old_input_details = old_model.get_input_details()
            old_output_details = old_model.get_output_details()
            st.write(f"Input shape: {old_input_details[0]['shape']}")
            st.write(f"Output shape: {old_output_details[0]['shape']}")
    
    with col2:
        st.subheader("New Model")
        new_model = load_model(new_model_path)
        if new_model:
            st.success("New model loaded successfully")
            new_input_details = new_model.get_input_details()
            new_output_details = new_model.get_output_details()
            st.write(f"Input shape: {new_input_details[0]['shape']}")
            st.write(f"Output shape: {new_output_details[0]['shape']}")
    
    # Load sentences
    sentences = load_sentences(sentences_path)
    st.header("Sentences")
    st.write(f"Loaded {len(sentences)} sentences")
    if st.checkbox("Show loaded sentences"):
        st.write(sentences[:10])
        if len(sentences) > 10:
            st.write("...")
    
    # Pre-compute embeddings for all sentences (do this only once for efficiency)
    if 'old_sentence_embeddings' not in st.session_state or st.button("Recompute Embeddings"):
        st.session_state.old_sentence_embeddings = []
        st.session_state.new_sentence_embeddings = []
        
        if old_model and new_model:
            progress_bar = st.progress(0)
            st.write("Computing sentence embeddings...")
            
            for i, sentence in enumerate(sentences):
                if i % 10 == 0:
                    progress_bar.progress(i / len(sentences))
                
                old_embedding, _ = get_embedding(sentence, old_model, tokenizer)
                new_embedding, _ = get_embedding(sentence, new_model, tokenizer)
                
                if old_embedding is not None:
                    st.session_state.old_sentence_embeddings.append(old_embedding[0])
                
                if new_embedding is not None:
                    st.session_state.new_sentence_embeddings.append(new_embedding[0])
            
            progress_bar.progress(1.0)
            st.write("Embeddings computed!")
    
    # Search interface
    st.header("Search")
    query = st.text_input("Enter a search query")
    
    if query and old_model and new_model:
        # Display tokenization for the query (for debugging)
        with st.expander("View tokenization"):
            tokens = tokenizer.encode(query, out_type=int)
            pieces = tokenizer.encode(query, out_type=str)
            st.write("Token IDs:", tokens)
            st.write("Token pieces:", pieces)
        
        # Get query embeddings
        old_query_embedding, old_time = get_embedding(query, old_model, tokenizer)
        new_query_embedding, new_time = get_embedding(query, new_model, tokenizer)
        
        # Find similar sentences
        old_results = find_similar_sentences(
            old_query_embedding, 
            st.session_state.old_sentence_embeddings,
            sentences
        )
        
        new_results = find_similar_sentences(
            new_query_embedding, 
            st.session_state.new_sentence_embeddings,
            sentences
        )
        
        # Add rank information
        for i, result in enumerate(old_results):
            result["rank"] = i + 1
            
        for i, result in enumerate(new_results):
            result["rank"] = i + 1
        
        # Create separate dataframes
        old_df = pd.DataFrame([
            {"Sentence": r["sentence"], "Similarity": f"{r['similarity']:.4f}", "Rank": r["rank"]} 
            for r in old_results
        ])
        
        new_df = pd.DataFrame([
            {"Sentence": r["sentence"], "Similarity": f"{r['similarity']:.4f}", "Rank": r["rank"]} 
            for r in new_results
        ])
        
        # Display results in two columns
        st.subheader("Search Results")
        col1, col2 = st.columns(2)
        
        with col1:
            st.markdown("### Old Model Results")
            st.dataframe(old_df, use_container_width=True)
        
        with col2:
            st.markdown("### New Model Results")
            st.dataframe(new_df, use_container_width=True)
        
        # Show timing information
        st.subheader("Inference Time")
        st.write(f"Old model: {old_time * 1000:.2f} ms")
        st.write(f"New model: {new_time * 1000:.2f} ms")
        st.write(f"Speed improvement: {old_time / new_time:.2f}x")
        
        # Show embedding visualizations
        st.subheader("Embedding Visualizations")
        col1, col2 = st.columns(2)
        
        with col1:
            st.write("Old Model Embedding (first 20 dimensions)")
            st.bar_chart(pd.DataFrame({
                'value': old_query_embedding[0][:20]
            }))
        
        with col2:
            st.write("New Model Embedding (first 20 dimensions)")
            st.bar_chart(pd.DataFrame({
                'value': new_query_embedding[0][:20]
            }))

if __name__ == "__main__":
    main()