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.")