missantroop's picture
Update app.py
089ba1a verified
import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
@st.cache_resource
def load_model():
model_path = "missantroop/CorebotClassifier" # Substitua pelo caminho do seu modelo
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(
model_path, torch_dtype=torch.float32, device_map="cpu"
)
return tokenizer, model
# Função para classificar texto
def classify_lyrics(model, tokenizer, text):
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
outputs = model.generate(**inputs, max_new_tokens=50)
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return generated_text
# Interface Streamlit
st.title("Classificador de Letras de Metal")
st.write("Carregue uma letra de música para classificação.")
# Input do usuário
lyrics = st.text_area("Insira a letra da música aqui", height=200)
if st.button("Classificar"):
if lyrics.strip():
st.write("Carregando modelo...")
tokenizer, model = load_model()
st.write("Modelo carregado! Classificando...")
try:
result = classify_lyrics(model, tokenizer, lyrics)
st.success("Classificação concluída!")
st.write(f"**Resultado da classificação:** {result}")
except Exception as e:
st.error(f"Erro ao classificar: {str(e)}")
else:
st.warning("Por favor, insira uma letra antes de classificar.")