Spaces:
Sleeping
Sleeping
File size: 1,208 Bytes
5966291 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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}")
|