Update handler.py

handler.py

@@ -2,7 +2,6 @@ import os
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Global variables for model, tokenizer, and device
 model = None
 tokenizer = None
 device = None
@@ -13,44 +12,56 @@ def init():
     """
     global model, tokenizer, device
     
-    # Set your model name or path here
+    # Replace this with your model repository ID
     model_name_or_path = "0xroyce/NazareAI-Senior-Marketing-Strategist"
     
+    # Load the tokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+
+    # Load the model
     model = AutoModelForCausalLM.from_pretrained(
         model_name_or_path,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         low_cpu_mem_usage=True
     )
     
+    # Set up the device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
-    model.eval()  # Put model in evaluation mode
+    model.eval()
+
+    # Store in global variables
+    globals()["model"] = model
+    globals()["tokenizer"] = tokenizer
+    globals()["device"] = device
+
 
 def inference(model_inputs: dict) -> dict:
     """
-    This function is called for every request and should return the model's output.
-    The input is a dictionary and the output should be a dictionary.
+    This function is called for every request.
+    The input is a dictionary with a 'prompt' key.
+    The output is a dictionary with 'generated_text'.
     """
     global model, tokenizer, device
 
-    # Extract the prompt from the input
+    # Get the prompt from the input
     prompt = model_inputs.get("prompt", "")
     if not prompt:
         return {"error": "No prompt provided."}
 
-    # Tokenize inputs
+    # Tokenize the prompt
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
-    # Generate text
-    # You can adjust parameters like max_new_tokens, temperature, or top_p as needed
+    # Run generation
     output_ids = model.generate(
-        **inputs, 
-        max_new_tokens=200, 
-        do_sample=True, 
-        top_p=0.9, 
+        **inputs,
+        max_new_tokens=200,
+        do_sample=True,
+        top_p=0.9,
         temperature=0.7
     )
+
+    # Decode the output
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
 
     return {"generated_text": output_text}