Update handler

README.md

@@ -26,10 +26,10 @@ from peft import PeftConfig, PeftModel
 
 # Import the model
 config = PeftConfig.from_pretrained("JeremyArancio/llm-tolkien")
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map='auto')
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_8bit=True, device_map='auto')
 tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 # Load the Lora model
-model = PeftModel.from_pretrained(model, hf_repo)
+model = PeftModel.from_pretrained(model, "JeremyArancio/llm-tolkien")
 ```
 
 # Run the model

handler.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Any
+from typing import Dict, Any
 import logging
 
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -6,6 +6,7 @@ from peft import PeftConfig, PeftModel
 
 
 LOGGER = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
 
 
 class EndpointHandler():
@@ -16,26 +17,24 @@ class EndpointHandler():
         # Load the Lora model
         self.model = PeftModel.from_pretrained(model, path)
 
-    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """
         Args:
             data (Dict): The payload with the text prompt and generation parameters.
         """
         LOGGER.info(f"Received data: {data}")
         # Get inputs
-        inputs = data.pop("inputs", data)
+        prompt = data.pop("prompt", data)
         parameters = data.pop("parameters", None)
-        LOGGER.info("Data extracted.")
         # Preprocess
-        LOGGER.info(f"Start tokenizer: {inputs}")
-        inputs_ids = self.tokenizer(inputs, return_tensors="pt").inputs_ids
+        input = self.tokenizer(prompt, return_tensors="pt")
         # Forward
         LOGGER.info(f"Start generation.")
         if parameters is not None:
-            outputs = self.model.generate(inputs_ids, **parameters)
+            output = self.model.generate(**input, **parameters)
         else:
-            outputs = self.model.generate(inputs_ids)
+            output = self.model.generate(**input)
         # Postprocess
-        prediction = self.tokenizer.decode(outputs[0])
+        prediction = self.tokenizer.decode(output[0])
         LOGGER.info(f"Generated text: {prediction}")
         return {"generated_text": prediction}