Update app.py

app.py

@@ -28,14 +28,11 @@ model_id = "fireworks-ai/firefunction-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",  # Permitir el offload automático entre CPU y GPU
+    device_map="auto",  # Permitir offloading automático entre CPU y GPU
     torch_dtype=torch.float16,
     quantization_config=quant_config
 )
 
-# Asegurar que el modelo esté en la GPU
-model.to("cuda")
-
 # Establecer el token de relleno
 if tokenizer.pad_token_id is None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
@@ -68,12 +65,19 @@ if uploaded_file is not None:
         batch_size = 4  # Reducir batch size para minimizar el uso de memoria
         job_titles_batches = [job_titles[i:i+batch_size] for i in range(0, len(job_titles), batch_size)]
 
-        # Definir el prompt para Fireworks
+        # Definir el prompt para Fireworks con formato de función
         prompt_template = PromptTemplate(
             template=(
-                "Calculate the cosine similarity between the query: '{query}' "
-                "and the list of job titles: {job_titles}. "
-                "Return the results as 'Job Title: [Job Title], Score: [Cosine Similarity Score]'."
+                "Function: calculate_cosine_similarity\n"
+                "Description: Calculate the cosine similarity between the given query and job titles.\n"
+                "Parameters:\n"
+                "  - query: The query string to compare.\n"
+                "  - job_titles: A list of job titles to compare against.\n"
+                "Input:\n"
+                "  query: '{query}'\n"
+                "  job_titles: {job_titles}\n"
+                "Output:\n"
+                "  Return the results as 'Job Title: [Job Title], Score: [Cosine Similarity Score]'."
             ),
             input_variables=["query", "job_titles"]
         )
@@ -90,13 +94,13 @@ if uploaded_file is not None:
                 all_scores = []
                 try:
                     for batch in job_titles_batches:
-                        # Tokenizar la entrada y mover a CUDA
+                        # Tokenizar la entrada y mover los tensores a CUDA
                         model_inputs = tokenizer(
                             batch,
                             return_tensors="pt",
                             padding=True,
                             truncation=True
-                        ).to("cuda")  # Mover explícitamente a CUDA
+                        ).to("cuda")  # Mover solo los tensores de entrada a CUDA
 
                         with torch.cuda.amp.autocast():  # Usar Mixed Precision
                             model_inputs['attention_mask'] = (model_inputs['input_ids'] != tokenizer.pad_token_id).int().to("cuda")