import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# Load pre-trained model and tokenizer
model_name = "distilbert-base-uncased"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Define a function to search for Samba songs
def search_samba(query):
    # Tokenize the query
    inputs = tokenizer(query, return_tensors="pt")

    # Get the model's output
    outputs = model(**inputs)

    # Get the predicted class probabilities
    logits = outputs.logits
    probs = torch.nn.functional.softmax(logits, dim=1)

    # Return the top 5 Samba song matches
    top_5 = torch.topk(probs, k=5)
    return top_5.indices, top_5.values

# Create a Streamlit app
st.title("Samba Search")

# Get user input
query = st.text_input("Enter a song title or artist")

# Search for Samba songs
if st.button("Search"):
    indices, values = search_samba(query)

    # Display the top 5 matches
    st.write("Top 5 Samba Song Matches:")
    for i, (index, value) in enumerate(zip(indices[0], values[0])):
        st.write(f"{i+1}. Song {index.item()} - Probability: {value.item():.2f}")