handler.py

import os
from typing import Dict, List, Any
import torch
from unsloth import FastLanguageModel

class EndpointHandler:
    def __init__(self, model_dir: str = ""):
        print(f"[DEBUG] Original model_dir: {model_dir}")
        # Override model_dir with your checkpoint path.
        model_dir = "RichardLu/mistral7b_aspectsentiment_laptop"
        print(f"[DEBUG] Using model_dir: {model_dir}")
        
        # Inference configuration.
        max_seq_length = 2048
        load_in_4bit = True  # or False as needed.
        
        # Load the model and tokenizer.
        self.model, self.tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_dir,
            max_seq_length=max_seq_length,
            load_in_4bit=load_in_4bit
        )
        print("[DEBUG] Model and tokenizer loaded successfully.")
        
        # Set the model to inference mode.
        FastLanguageModel.for_inference(self.model)
        
        # Define the instructabsa instruction text with examples.
        self.instructabsa_instruction = """Definition: The output will be 'positive' if the aspect identified in the sentence contains a positive sentiment. If the sentiment of the identified aspect in the input is negative the answer will be 'negative'.
Otherwise, the output should be 'neutral'. For aspects which are classified as noaspectterm, the sentiment is none.
Positive example 1-
input: I charge it at night and skip taking the cord with me because of the good battery life. The aspect is battery life.
output: positive
Positive example 2-
input: I even got my teenage son one, because of the features that it offers, like, iChat, Photobooth, garage band and more!. The aspect is garage band.
output: positive
Negative example 1-
input: Speaking of the browser, it too has problems. The aspect is browser.
output: negative
Negative example 2-
input: The keyboard is too slick. The aspect is keyboard.
output: negative
Neutral example 1-
input: I took it back for an Asus and same thing- blue screen which required me to remove the battery to reset. The aspect is battery.
output: neutral
Neutral example 2-
input: Nightly my computer defrags itself and runs a virus scan. The aspect is virus scan.
output: neutral
Now complete the following example-
"""
        
        # Define the Alpaca-style prompt template.
        self.alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
### Instruction:
{}
### Input:
{}
### Response:
{}"""

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        # First, try to use the "inputs" key.
        if "inputs" in data:
            full_input = data["inputs"]
        # Otherwise, check for "review" and "aspect" keys.
        elif "review" in data and "aspect" in data:
            full_input = f"{data['review']} The aspect is {data['aspect']}."
        else:
            return [{"error": "Provide either an 'inputs' key or both 'review' and 'aspect' keys."}]
        
        # Build the final prompt.
        prompt = self.alpaca_prompt.format(self.instructabsa_instruction, full_input, "")
        
        # Set device.
        device = "cuda" if torch.cuda.is_available() else "cpu"
        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(device)
        
        # Generate the output.
        output_ids = self.model.generate(**inputs, max_new_tokens=128)
        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
        
        # Extract the predicted sentiment.
        if "### Response:" in output_text:
            predicted = output_text.split("### Response:")[-1].strip()
        else:
            predicted = output_text.strip()
        
        return [{"predicted": predicted}]