add support for phi

main.py

@@ -16,8 +16,11 @@ def parse_arguments():
     parser = argparse.ArgumentParser()
     parser.add_argument('--whisper_tensorrt_path',
                         type=str,
-                        default=None,
+                        default="/root/TensorRT-LLM/examples/whisper/whisper_small_en",
                         help='Whisper TensorRT model path')
+    parser.add_argument('--mistral',
+                        action="store_true",
+                        help='Mistral')
     parser.add_argument('--mistral_tensorrt_path',
                         type=str,
                         default=None,
@@ -26,6 +29,17 @@ def parse_arguments():
                         type=str,
                         default="teknium/OpenHermes-2.5-Mistral-7B",
                         help='Mistral TensorRT model path')
+    parser.add_argument('--phi',
+                        action="store_true",
+                        help='Phi')
+    parser.add_argument('--phi_tensorrt_path',
+                        type=str,
+                        default="/root/TensorRT-LLM/examples/phi/phi_engine",
+                        help='Phi TensorRT model path')
+    parser.add_argument('--phi_tokenizer_path',
+                        type=str,
+                        default="/root/TensorRT-LLM/examples/phi/phi-2",
+                        help='Phi Tokenizer path')
     return parser.parse_args()
 
 
@@ -36,10 +50,17 @@ if __name__ == "__main__":
         import sys
         sys.exit(0)
     
-    if not args.mistral_tensorrt_path or not args.mistral_tokenizer_path:
-        raise ValueError("Please provide mistral_tensorrt_path and mistral_tokenizer_path to run the pipeline.")
-        import sys
-        sys.exit(0)
+    if args.mistral:
+        if not args.mistral_tensorrt_path or not args.mistral_tokenizer_path:
+            raise ValueError("Please provide mistral_tensorrt_path and mistral_tokenizer_path to run the pipeline.")
+            import sys
+            sys.exit(0)
+
+    if args.phi:
+        if not args.phi_tensorrt_path or not args.phi_tokenizer_path:
+            raise ValueError("Please provide phi_tensorrt_path and phi_tokenizer_path to run the pipeline.")
+            import sys
+            sys.exit(0)
 
     multiprocessing.set_start_method('spawn')
     
@@ -70,8 +91,10 @@ if __name__ == "__main__":
     llm_process = multiprocessing.Process(
         target=llm_provider.run,
         args=(
-            args.mistral_tensorrt_path,
-            args.mistral_tokenizer_path,
+            # args.mistral_tensorrt_path,
+            # args.mistral_tokenizer_path,
+            args.phi_tensorrt_path,
+            args.phi_tokenizer_path,
             transcription_queue,
             llm_queue,
         )