Update vid2persona/init.py

vid2persona/init.py

@@ -14,6 +14,14 @@ ALLOWED_LLM_FOR_HF_PRO_ACCOUNTS = [
     # "openchat/openchat-3.5-0106"
 ]
 
+def init_model(model_id):
+    global tokenizer
+    global model
+    
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    tokenizer.use_default_system_prompt = False
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
+
 def auth_gcp():
     gcp_credentials = os.getenv("GCP_CREDENTIALS")
     with open("gcp-credentials.json", "w") as f: