import os
import shutil
from transformers import LlamaTokenizer, LlamaForCausalLM
import torch

# Definir o diretório de cache para o modelo
os.environ['TRANSFORMERS_CACHE'] = '/tmp/model_cache'

# Limpar o diretório de cache, se existir
cache_dir = '/tmp/model_cache'
if os.path.exists(cache_dir):
    shutil.rmtree(cache_dir)

def carregar_modelo():
    tokenizer = LlamaTokenizer.from_pretrained("lzw1008/Emollama-7b")
    model = LlamaForCausalLM.from_pretrained("lzw1008/Emollama-7b", device_map="auto")
    return tokenizer, model

tokenizer, model = carregar_modelo()

def analise_sentimento(texto: str) -> str:
    prompt = f"""Humano:
Tarefa: Categorize a emoção expressada no texto como 'neutro ou sem emoção' ou identifique a presença de uma ou mais emoções das listadas (raiva, confusão, ansiedade, desgosto, medo, alegria, amor, otimismo, pessimismo, tristeza, surpresa, confiança).
Texto: {texto}
Esse texto contém a emoção:

Assistente:"""

    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(model.device)

    with torch.no_grad():
        generate_ids = model.generate(
            inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            max_length=256,
            do_sample=True,
            top_p=0.95,
            temperature=0.7,
            pad_token_id=tokenizer.eos_token_id
        )

    resposta = tokenizer.decode(generate_ids[0], skip_special_tokens=True)
    
    return resposta