import streamlit as st
import pandas as pd
import torch
from langchain import LLMChain, PromptTemplate
from langchain.llms import HuggingFaceHub
from sentence_transformers import SentenceTransformer, util

# Configuración de la app
st.title("Subida de archivo CSV e interacción con Llama 3.1")

# Subir archivo CSV
uploaded_file = st.file_uploader("Sube un archivo CSV", type=["csv"])

if uploaded_file is not None:
    # Leer el archivo CSV en un DataFrame
    df = pd.read_csv(uploaded_file)

    # Mostrar una vista previa del archivo
    st.write("Vista previa del archivo CSV:")
    st.write(df.head())
    
    # Especificar la columna de títulos de trabajo
    job_title_column = st.selectbox("Selecciona la columna de títulos de trabajo", df.columns)

    # Definir el prompt
    prompt = "Toma el query 'aspiring human resources specialist' y calcula el cosine similarity score de cada job_title. Crea una nueva columna 'Score' con dichos valores."

    # Procesar el prompt con Langchain cuando se presione el botón
    if st.button("Calcular Similitud"):
        # Paso 1: Usar Llama para generar texto o aclaraciones
        prompt_template = PromptTemplate(
            template="El usuario ha subido un CSV con la columna de títulos de trabajo. {user_prompt}",
            input_variables=["user_prompt"]
        )

        # Instancia del modelo Llama en HuggingFace
        hf_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
        llm = HuggingFaceHub(repo_id="meta-llama/LLaMA-3.1-8B-Instruct", huggingfacehub_api_token=hf_token)
        llm_chain = LLMChain(llm=llm, prompt=prompt_template)

        # Ejecutar Llama
        llama_response = llm_chain.run({
            "user_prompt": prompt
        })

        # Mostrar la respuesta generada por Llama
        st.write("Respuesta generada por Llama:")
        st.write(llama_response)
        
        # Paso 2: Usar Sentence Transformers para calcular la similitud de coseno
        model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

        # Query que queremos comparar
        query = "aspiring human resources specialist"

        # Obtener embeddings del query
        query_embedding = model.encode(query, convert_to_tensor=True)

        # Obtener embeddings de los títulos de trabajo
        job_titles = df[job_title_column].astype(str).tolist()
        job_title_embeddings = model.encode(job_titles, convert_to_tensor=True)

        # Calcular la similitud de coseno
        cosine_scores = util.pytorch_cos_sim(query_embedding, job_title_embeddings)

        # Agregar los puntajes de similitud al DataFrame original
        df['Score'] = cosine_scores.squeeze().tolist()

        # Mostrar el DataFrame actualizado
        st.write("DataFrame con los puntajes de similitud:")
        st.write(df)
else:
    st.write("Por favor, sube un archivo CSV.")