Update main_process/moe_router.py

main_process/moe_router.py

@@ -61,18 +61,30 @@ class NState(dict):
 
 # ---------------- LLM utilizado para el free_narration ----------------
 class SalamandraClient:
-    '''
-    def __init__(self, model_id="BSC-LT/salamandra-7b-instruct"):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+    def __init__(self, model_id="BSC-LT/salamandra-7b-instruct-tools"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
         self.model = AutoModelForCausalLM.from_pretrained(
             model_id,
             device_map="auto",
-            torch_dtype=torch.bfloat16
+            torch_dtype=torch.bfloat16,
+            use_auth_token= HF_TOKEN
         )
-    '''
 
-    def chat(self, prompt) -> str:
-        return get_from_prompt(prompt)
+    def chat(self, message, tools) -> str:
+        date_string = datetime.today().strftime('%Y-%m-%d')
+
+        prompt = self.tokenizer.apply_chat_template(
+            message,
+            tokenize=False,
+            add_generation_prompt=True,
+            date_string=date_string,
+            tools=tools
+        )
+
+        inputs = self.tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+        outputs = self.model.generate(input_ids=inputs.to(self.model.device), max_new_tokens=200)
+
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Esto aquí sólo se utiliza para la valoración:
 class GPT5Client:
@@ -1469,7 +1481,7 @@ async def generate_salamadra_result(
                     state = tool(state, srt_final, csv_evaluacion)
                 
                 elif isinstance(tool, Free_Narration):
-                    state = tool(state, srt_final, free_narration_moe)
+                    state = tool(state, srt_final, free_narration_salamandra)
     
                 messages_registro.append({"role": "assistant", "content": f"Ejecuté {tool_name} correctamente."})
         else:
@@ -1479,64 +1491,82 @@ async def generate_salamadra_result(
     
     messages_registro = [{"role": "system", "content": "Eres un agente que puede ejecutar herramientas Python usando las herramientas disponibles."}]
     
-    count = 1
+    import yaml
 
-    user_prompt = "Ejecuta la función add_ad"
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, state, tools, user_prompt, messages_registro, count)
-    state = add_ad(state, temp_srt_files[0].name, temp_srt_files[1].name)
-    print("Transcripción con AD guardada")
+    with open("config.yaml", "r") as file:
+        config = yaml.safe_load(file)
     
-    user_prompt = "Ejecuta la función add_silence_ad"
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    state = add_silence_ad(state, temp_srt_files[1].name, temp_srt_files[2].name)
-    print("Transcripción con AD y Add_Silence_AD guardada")
+    stools_enabled = config.get("stools", False) 
+
+    if stools_enabled:
+        count = 1
+        user_prompt = "Ejecuta la función add_ad"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función add_silence_ad"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función unir_ad_silence"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        count = 2
+        user_prompt = "Ejecuta la función add_silence_ad"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función Unir_AD_Silences_a_ADs"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función Introduccion_OCR"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+ 
+        user_prompt = "Ejecuta la función Identity_Manager"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función UNE_Actor_prev"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función UNE_Actor"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función Valoracion_Final"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+        user_prompt = "Ejecuta la función Free_Narration"
+        final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
+
+    else:
+        state = add_ad(state, temp_srt_files[0].name, temp_srt_files[1].name)
+        print("Transcripción con AD guardada")
     
-    user_prompt = "Ejecuta la función unir_ad_silence"
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    state = unir_ad_silence(state, temp_srt_files[2].name, temp_srt_files[3].name)
-    print("Transcripción con AD y Unir_AD_Silence guardada")
+        state = add_silence_ad(state, temp_srt_files[1].name, temp_srt_files[2].name)
+        print("Transcripción con AD y Add_Silence_AD guardada")
     
-    count = 2
+        state = unir_ad_silence(state, temp_srt_files[2].name, temp_srt_files[3].name)
+        print("Transcripción con AD y Unir_AD_Silence guardada")
     
-    user_prompt = "Ejecuta la función add_silence_ad"
-    state = add_silence_ad(state, temp_srt_files[3].name, temp_srt_files[4].name)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD y Add_Silence_AD guardada")
+        state = add_silence_ad(state, temp_srt_files[3].name, temp_srt_files[4].name)
+        print("Transcripción con AD y Add_Silence_AD guardada")
     
-    user_prompt = "Ejecuta la función Unir_AD_Silences_a_ADs"
-    state = unir_ad_silences_a_ads(state, temp_srt_files[4].name, temp_srt_files[5].name)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD y Unir_AD_Silences_a_ADs guardada")
+        state = unir_ad_silences_a_ads(state, temp_srt_files[4].name, temp_srt_files[5].name)
+        print("Transcripción con AD y Unir_AD_Silences_a_ADs guardada")
     
-    user_prompt = "Ejecuta la función Introduccion_OCR"
-    state = introduccion_ocr(state, temp_srt_files[5].name, temp_srt_files[6].name)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD, Add_Silence_AD e Introduccion_OCR guardada")
+        state = introduccion_ocr(state, temp_srt_files[5].name, temp_srt_files[6].name)
+        print("Transcripción con AD, Add_Silence_AD e Introduccion_OCR guardada")
     
-    user_prompt = "Ejecuta la función Identity_Manager"
-    state = identity_manager(state, temp_srt_files[6].name, temp_srt_files[7].name)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR e Identity_Manager guardada")
+        state = identity_manager(state, temp_srt_files[6].name, temp_srt_files[7].name)
+        print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR e Identity_Manager guardada")
     
-    user_prompt = "Ejecuta la función UNE_Actor_prev"
-    state = une_actor_prev(state, temp_srt_files[7].name, temp_srt_files[8].name)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR, Identity_Manager y norma UNE guardada")
+        state = une_actor_prev(state, temp_srt_files[7].name, temp_srt_files[8].name)
+        print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR, Identity_Manager y norma UNE guardada")
     
-    user_prompt = "Ejecuta la función UNE_Actor"
-    state = une_actor(state, temp_srt_files[8].name, srt_final)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR, Identity_Manager y norma UNE guardada")
+        state = une_actor(state, temp_srt_files[8].name, srt_final)
+        print("Transcripción con AD, Add_Silence_AD, Introduccion_OCR, Identity_Manager y norma UNE guardada")
     
-    user_prompt = "Ejecuta la función Valoracion_Final"
-    state = valoracion_final(state, srt_final, csv_evaluacion)
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    print("Valoración guardada")
+        state = valoracion_final(state, srt_final, csv_evaluacion)
+        print("Valoración guardada")
     
-    user_prompt = "Ejecuta la función Free_Narration"
-    #final_state, messages_registro = run_salamandra_agent(salamandraclient, final_state, tools, user_prompt, messages_registro, count)
-    state = free_narration(state, srt_final, free_narration_salamandra)
-    print("Free Narration guardada")
+        state = free_narration(state, srt_final, free_narration_salamandra)
+        print("Free Narration guardada")
 
     srt_update(srt_final,srt_final)