import torch
import pandas as pd
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
from langchain.llms import HuggingFacePipeline
from huggingface_hub import login
from pydantic import BaseModel, model_validator

# Token secreto de Hugging Face
huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
login(huggingface_token)

# Cargar el modelo Llama 3.1
model_name = "meta-llama/llama-3.1-8b-instruct"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

pipe = pipeline("feature-extraction", model=model, tokenizer=tokenizer)

# Usar transformers pipeline para cargar el modelo y tokenizer
llm_pipeline = HuggingFacePipeline(pipeline=pipe)

# Interfaz de Streamlit
st.title("Cosine Similarity con Llama 3.1")

# Subida de archivo CSV
uploaded_file = st.file_uploader("Sube un archivo CSV", type=["csv"])
if uploaded_file is not None:
    # Leer el archivo CSV
    df = pd.read_csv(uploaded_file)

    # Solicitar un query al usuario
    query = st.text_input("Por favor ingresa tu consulta:")
    
    if query:
        # Crear el prompt para calcular la similitud del coseno
        prompt = f"""
        Con la columna 'job_title' del csv, tiene que calcular los cosine similarity score entre cada registro y el query,
        tiene que crear una columna nueva 'Score' con dichos scores y además tiene que ordenar los registros de mayor a menor valor de 'Score'.
        El query es: {query}
        """

        # Llamar al modelo con el prompt
        try:
            response = llm_pipeline(prompt)
            # Mostrar la respuesta del modelo
            st.write("Respuesta del modelo:", response)

        except Exception as e:
            st.error(f"Ocurrió un error al procesar el modelo: {str(e)}")

# Modelo de validación de datos con Pydantic
class ConsultaModelo(BaseModel):
    query: str

    @model_validator(mode='before')
    def validar_query(cls, values):
        query = values.get('query')
        if not query:
            raise ValueError("El query no puede estar vacío.")
        return values

# Ejemplo de uso del modelo de validación
if query:
    try:
        valid_query = ConsultaModelo(query=query)
        st.success("Consulta válida.")
    except ValueError as e:
        st.error(f"Error de validación: {e}")