Update handler.py

handler.py

@@ -9,20 +9,23 @@ class EndpointHandler:
         # Retrieve the Hugging Face token from the environment variable.
         hf_token = os.getenv("HF_TOKEN", None)
         
-        # Configuration for inference
+        # If the endpoint provides '/repository', override with the intended model ID.
+        if model_dir == "/repository" or model_dir.strip() == "":
+            model_dir = "RichardLu/Mistral7b_AE_res"
+        
+        # Configuration for inference.
         max_seq_length = 2048
         dtype = None  # Auto-detect data type; adjust if needed.
         load_in_4bit = True  # Use 4-bit quantization to reduce memory usage.
         
         # Load the model and tokenizer.
-        # If a model_dir is provided by the endpoint, it will load from that path;
-        # otherwise, it uses a default UnsLoth Mistral model.
         self.model, self.tokenizer = FastLanguageModel.from_pretrained(
-            model_name=RichardLu/Mistral7b_AE_res,
+            model_name=model_dir,
             max_seq_length=max_seq_length,
             dtype=dtype,
             load_in_4bit=load_in_4bit,
-            token=hf_token  # Pass the HF token if provided.
+            token=hf_token,            # Use the HF token if provided.
+            trust_remote_code=True      # Enable if your model requires remote code.
         )
         
         # Switch the model to inference mode.
@@ -30,7 +33,8 @@ class EndpointHandler:
         
         # Define the instruction text with examples for aspect extraction.
         self.instructabsa_instruction = (
-            "Definition: The output will be the aspects (both implicit and explicit) which have an associated opinion that are extracted from the input text. In cases where there are no aspects the output should be noaspectterm.\n"
+            "Definition: The output will be the aspects (both implicit and explicit) which have an associated opinion that are extracted from the input text. "
+            "In cases where there are no aspects the output should be noaspectterm.\n"
             "Positive example 1-\n"
             "input: With the great variety on the menu, I eat here often and never get bored.\n"
             "output: menu\n"