Update app.py

app.py

@@ -1,105 +1,67 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import pandas as pd
 import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from langchain.llms import HuggingFacePipeline
 from huggingface_hub import login
-import pandas as pd
-from threading import Thread
 
-# Token Secret de Hugging Face
+# Token de Hugging Face (Secreto)
 huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
 login(huggingface_token)
 
-# Cargar el tokenizador y el modelo
-model_id = "meta-llama/Llama-3.2-1B"
+# Cargar el modelo Llama 3.1 y el tokenizador
+model_id = "meta-llama/Llama-3.1-1B"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
-tokenizer.pad_token = tokenizer.eos_token
-
-MAX_INPUT_TOKEN_LENGTH = 10000
 
-def generate_response(input_text, temperature=0.7, max_new_tokens=20):
-    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
+# Crear pipeline de generación de texto
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=1024)
+llm_pipeline = HuggingFacePipeline(pipeline=pipe)
 
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        st.warning(f"Se recortó la entrada porque excedió el límite de {MAX_INPUT_TOKEN_LENGTH} tokens.")
+# Interfaz de Streamlit
+st.title("Cosine Similarity with Llama 3.1")
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=120.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        input_ids=input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_k=20,
-        top_p=0.9,
-        temperature=temperature,
-        num_return_sequences=3,
-        eos_token_id=[tokenizer.eos_token_id]
-    )
+# Subir archivo CSV
+uploaded_file = st.file_uploader("Sube un archivo CSV con la columna 'job_title':", type=["csv"])
 
-    try:
-        t = Thread(target=model.generate, kwargs=generate_kwargs)
-        t.start()
-        t.join()  # Asegura que la generación haya terminado
-
-        outputs = []
-        for text in streamer:
-            outputs.append(text)
-        if not outputs:
-            raise ValueError("No se generó ninguna respuesta.")
-        
-        response = "".join(outputs).strip().split("\n")[0]
-        return response
-    except Exception as e:
-        st.error(f"Error durante la generación: {e}")
-        return "Error en la generación de texto."
-
-def main():
-    st.title("Chat con Meta Llama 3.2 1B")
-    
-    uploaded_file = st.file_uploader("Por favor, sube un archivo CSV para iniciar:", type=["csv"])
+if uploaded_file is not None:
+    # Cargar el CSV en un DataFrame
+    df = pd.read_csv(uploaded_file)
     
-    if uploaded_file is not None:
-        df = pd.read_csv(uploaded_file)
+    if 'job_title' in df.columns:
         query = 'aspiring human resources specialist'
-        value = 0.00
-        if 'job_title' in df.columns:
-            job_titles = df['job_title']
+        job_titles = df['job_title'].tolist()
 
-            # Definir el prompt con in-context learning
-            initial_prompt = (
-                "Step 1: Extract the first record from the dataframe df.\n"
-                f" {df.iloc[0]['job_title']}\n"
-                #f"List: {job_titles}\n"
-                #"First job title: \n"
-                #"\n"
-                "Step 2: Calculate the cosine similarity score between the job_title of the extracted record {df.iloc[0]['job_title']} and the given {query} and assign it to {value}.\n"
-                f"Query: '{query}'\n"
-                "Cosine similarity score: \n"
-                "Step 3: Print the value of the calculated cosine similarity"
-                f"Result: {value}"
-            )
+        # Definir el prompt para el LLM
+        prompt = (
+            f"You are given a query and a list of job titles. Your task is to calculate the cosine similarity "
+            f"between the query and each job title. The query is: '{query}'. For each job title, provide the similarity "
+            f"score as a new column in the dataframe, called 'Score'. Return the dataframe with job titles and scores.\n"
+            f"Job Titles: {job_titles}\n"
+            f"Output format:\n"
+            f"1. Job Title: [Job Title], Score: [Cosine Similarity Score]\n"
+            f"2. Job Title: [Job Title], Score: [Cosine Similarity Score]\n"
+            f"..."
+        )
 
-            
-            st.write("Prompt inicial con In-context Learning:\n")
-            st.write(initial_prompt)
-            st.write(query)
+        # Mostrar el prompt inicial
+        st.write("Prompt enviado al LLM:")
+        st.write(prompt)
 
-            if st.button("Generar respuesta"):
-                with st.spinner("Generando respuesta..."):
-                    response = generate_response(initial_prompt, temperature=0.5)
-                    if response:
-                        st.write(f"Respuesta del modelo: {response}")
-                    else:
-                        st.warning("No se pudo generar una respuesta.")
+        # Generar respuesta del LLM
+        if st.button("Generar puntajes de similitud"):
+            with st.spinner("Calculando similitudes con Llama 3.1..."):
+                try:
+                    response = llm_pipeline(prompt)
+                    st.write("Respuesta del modelo:")
+                    st.write(response)
 
-                st.success("La conversación ha terminado.")
-                
-                if st.button("Iniciar nueva conversación"):
-                    st.experimental_rerun()
-                elif st.button("Terminar"):
-                    st.stop()
-        else:
-            st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
+                    # Simular la asignación de puntajes en la columna 'Score' (ya que el modelo no ejecuta cálculos reales)
+                    df['Score'] = [0.95] * len(df)  # Este paso es solo ilustrativo
 
-if __name__ == "__main__":
-    main()
+                    # Mostrar el dataframe actualizado
+                    st.write("DataFrame con los puntajes de similitud:")
+                    st.write(df)
+                except Exception as e:
+                    st.error(f"Error durante la generación: {e}")
+    else:
+        st.error("La columna 'job_title' no se encuentra en el archivo CSV.")