Update handler.py

handler.py

@@ -5,7 +5,6 @@ import base64
 
 from pyannote.audio import Pipeline
 from transformers import pipeline, AutoModelForCausalLM
-from diarization_utils import diarize
 from huggingface_hub import HfApi
 from pydantic import ValidationError
 from starlette.exceptions import HTTPException
@@ -22,16 +21,6 @@ class EndpointHandler():
         logger.info(f"Using device: {device.type}")
         torch_dtype = torch.float32 if device.type == "cpu" else torch.float16
 
-        self.assistant_model = AutoModelForCausalLM.from_pretrained(
-            model_settings.assistant_model,
-            torch_dtype=torch_dtype,
-            low_cpu_mem_usage=True,
-            use_safetensors=True
-        ) if model_settings.assistant_model else None
-
-        if self.assistant_model:
-            self.assistant_model.to(device)
-
         self.asr_pipeline = pipeline(
             "automatic-speech-recognition",
             model=model_settings.asr_model,
@@ -39,18 +28,6 @@ class EndpointHandler():
             device=device
         )
 
-        if model_settings.diarization_model:
-            # diarization pipeline doesn't raise if there is no token
-            HfApi().whoami(model_settings.hf_token)
-            self.diarization_pipeline = Pipeline.from_pretrained(
-                checkpoint_path=model_settings.diarization_model,
-                use_auth_token=model_settings.hf_token,
-            )
-            self.diarization_pipeline.to(device)
-        else:
-            self.diarization_pipeline = None
-            
-    
     def __call__(self, inputs):
         file = inputs.pop("inputs")
         file = base64.b64decode(file)
@@ -65,8 +42,7 @@ class EndpointHandler():
 
         generate_kwargs = {
             "task": parameters.task, 
-            "language": parameters.language,
-            "assistant_model": self.assistant_model if parameters.assisted else None
+            "language": parameters.language
         }
 
         try:
@@ -81,23 +57,10 @@ class EndpointHandler():
             logger.error(f"ASR inference error: {str(e)}")
             raise HTTPException(status_code=400, detail=f"ASR inference error: {str(e)}")
         except Exception as e:
-            logger.error(f"Unknown error diring ASR inference: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Unknown error diring ASR inference: {str(e)}")
-
-        if self.diarization_pipeline:
-            try:
-                transcript = diarize(self.diarization_pipeline, file, parameters, asr_outputs)
-            except RuntimeError as e:
-                logger.error(f"Diarization inference error: {str(e)}")
-                raise HTTPException(status_code=400, detail=f"Diarization inference error: {str(e)}")
-            except Exception as e:
-                logger.error(f"Unknown error during diarization: {str(e)}")
-                raise HTTPException(status_code=500, detail=f"Unknown error during diarization: {str(e)}")
-        else:
-            transcript = []
+            logger.error(f"Unknown error during ASR inference: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Unknown error during ASR inference: {str(e)}")
 
         return {
-            "speakers": transcript,
             "chunks": asr_outputs["chunks"],
             "text": asr_outputs["text"],
         }