Update app.py

app.py

@@ -1,8 +1,8 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from accelerate import init_empty_weights, load_checkpoint_and_dispatch, dispatch_model, infer_auto_device_map
 import streamlit as st
 from huggingface_hub import login
 import pandas as pd
+from threading import Thread
 
 # Token Secret de Hugging Face
 huggingface_token = st.secrets["HUGGINGFACEHUB_API_TOKEN"]
@@ -11,23 +11,13 @@ login(huggingface_token)
 # Cargar el tokenizador y el modelo
 model_id = "meta-llama/Llama-3.2-1B"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id) #, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 tokenizer.pad_token = tokenizer.eos_token
 
 MAX_INPUT_TOKEN_LENGTH = 10000
 
-# Cargar el modelo con disk_offload
-with init_empty_weights():
-    model = AutoModelForCausalLM.from_config(model_id)
-
-device_map = infer_auto_device_map(model, max_memory={"disk": "2GiB"}, no_split_module_classes=["LlamaDecoderLayer"])
-model = load_checkpoint_and_dispatch(model, model_id, device_map=device_map, offload_folder="offload_dir")
-
-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("cpu")  # Usar 'cpu' para mantener la compatibilidad
+    input_ids = tokenizer.encode(input_text, return_tensors='pt').to(model.device)
 
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
@@ -43,14 +33,22 @@ def generate_response(input_text, temperature=0.7, max_new_tokens=20):
         top_p=0.9,
         temperature=temperature,
         num_return_sequences=3,
-        eos_token_id=tokenizer.eos_token_id
+        eos_token_id=[tokenizer.eos_token_id]
     )
 
     try:
-        outputs = model.generate(**generate_kwargs)
-
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-        return response.split("\n")[0]
+        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."
@@ -63,20 +61,28 @@ def main():
     if uploaded_file is not None:
         df = pd.read_csv(uploaded_file)
         query = 'aspiring human resources specialist'
-        
+        value = 0.00
         if 'job_title' in df.columns:
-            job_titles = df['job_title'].tolist()
+            job_titles = df['job_title']
 
             # Definir el prompt con in-context learning
             initial_prompt = (
-                f"Extract the first record from the dataframe df.\n"
-                f"First job title: '{df.iloc[0]['job_title']}'\n"
-                f"Calculate the cosine similarity between this job title and the query: '{query}'.\n"
-                "Print the cosine similarity score."
+                "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}"
             )
 
+            
             st.write("Prompt inicial con In-context Learning:\n")
             st.write(initial_prompt)
+            st.write(query)
 
             if st.button("Generar respuesta"):
                 with st.spinner("Generando respuesta..."):
@@ -96,4 +102,4 @@ def main():
             st.error("La columna 'job_title' no se encuentra en el archivo CSV.")
 
 if __name__ == "__main__":
-    main()
+    main()