import streamlit as st
import json
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
import os

# Recuperar el token de los secrets de Hugging Face Spaces
hf_token = os.getenv('hf_token')

# Configuración del modelo LLaMA
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto",
    quantization_config=bnb_config,
    token=hf_token
)

SYS_PROMPT = """You are an AI assistant tasked with analyzing the sentiment of a given text at the aspect level, specifically related to public security issues in Ecuador, such as crime, violence, law enforcement, or corruption. For each aspect mentioned in the text, classify the sentiment as 'Positive', 'Neutral', or 'Negative' and provide a clear justification for each aspect. Return the response **only** in JSON format and in Spanish, without additional information."""

def generate(formatted_prompt):
    messages = [{"role": "system", "content": SYS_PROMPT}, {"role": "user", "content": formatted_prompt}]

    input_ids = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    attention_mask = input_ids.ne(tokenizer.pad_token_id)

    outputs = model.generate(
        input_ids,
        max_new_tokens=1024,
        attention_mask=attention_mask,
        pad_token_id=tokenizer.pad_token_id
    )
    response = outputs[0][input_ids.shape[-1]:]
    return tokenizer.decode(response, skip_special_tokens=True)

def analyze_sentiment(text):
    formatted_prompt = f"Analyze the sentiment of the following text: {text}"
    llama_response = generate(formatted_prompt)

    # Intentar parsear la respuesta a JSON
    try:
        llama_data = json.loads(llama_response)
        return llama_data
    except json.JSONDecodeError:
        return {"error": "Failed to parse the response"}

# Interfaz de Streamlit
st.title("Análisis de Sentimiento con LLaMA")
input_text = st.text_area("Ingresa el texto a analizar:")

if st.button("Analizar"):
    if input_text:
        with st.spinner("Analizando..."):
            result = analyze_sentiment(input_text)
            st.json(result)
    else:
        st.warning("Por favor, ingresa un texto para analizar.")