sreejith8100 commited on
Commit
b7bae01
·
verified ·
1 Parent(s): 9915b66

Update handler.py

Browse files
Files changed (1) hide show
  1. handler.py +12 -1
handler.py CHANGED
@@ -18,8 +18,9 @@ class EndpointHandler:
18
 
19
  def load_model(self):
20
  model_name = "openbmb/MiniCPM-V-2_6"
21
- hf_token = os.getenv("HF_AUTH_TOKEN") # Set this as a secret in Hugging Face Space/Endpoint
22
 
 
23
  self.tokenizer = AutoTokenizer.from_pretrained(
24
  model_name, trust_remote_code=True, use_auth_token=hf_token
25
  )
@@ -74,3 +75,13 @@ class EndpointHandler:
74
  return {"output": output}
75
  except Exception as e:
76
  return {"error": f"Inference failed: {e}"}
 
 
 
 
 
 
 
 
 
 
 
18
 
19
  def load_model(self):
20
  model_name = "openbmb/MiniCPM-V-2_6"
21
+ hf_token = os.getenv("HF_AUTH_TOKEN") # Ensure this is set in your environment
22
 
23
+ # Load tokenizer and model with the Hugging Face API token
24
  self.tokenizer = AutoTokenizer.from_pretrained(
25
  model_name, trust_remote_code=True, use_auth_token=hf_token
26
  )
 
75
  return {"output": output}
76
  except Exception as e:
77
  return {"error": f"Inference failed: {e}"}
78
+
79
+ def __call__(self, data):
80
+ """
81
+ This makes the class callable as Hugging Face expects the handler to be callable.
82
+ """
83
+ return self.predict(data)
84
+
85
+
86
+ # Hugging Face looks for a callable handler
87
+ handler = EndpointHandler()