LlamaMeshClassifier.py

import gc
import os
import json
from LabelProcessor import LabelProcessor
from MeshRenderer import MeshRenderer
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch


class LlamaMeshClassifier:
    def __init__(self, model_name=None, device="cuda", max_new_tokens=128, max_input_tokens=131072):
        """Initialize a causal LM based classifier for meshes.

        Args:
            model_name (str): Path to a causal LM model compatible with the tokenizer.
            device (str): Device for running the model (e.g. "cuda").
            max_new_tokens (int): Maximum generation length; clipped to a safe limit.
            max_input_tokens (int): Maximum allowed input token length.
        """
        self.device = device
        model_name = model_name or os.path.normpath(
            os.path.join(os.path.dirname(__file__), "..", "models", "LLaMA-Mesh-model")
        )
        self.model = AutoModelForCausalLM.from_pretrained(model_name).to(self.device)
        self.model.eval()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.clip_text_token_limit = 45
        self.max_new_tokens = min(max_new_tokens, self.clip_text_token_limit)
        self.max_input_tokens = max_input_tokens
        self.terminators = [
            self.tokenizer.eos_token_id,
            self.tokenizer.convert_tokens_to_ids("<|eot_id|>")
        ]
        self.mesh_renderer = MeshRenderer()
        self.classification_prompt = """
            You are performing object classification.
            Output exactly ONE label containing 1 to 3 words. 
            The label must be the most specific identifiable real-world object category.
            If the object represents a well-known fictional character, brand, or copyrighted entity, output its proper name (e.g., 'Iron Man', 'Spider-Man', 'Batman'). 
            Do NOT describe shape, color, or geometry. 
            Do NOT output a sentence.
            Do NOT explain your reasoning.
            Return only the object name.
        """

    def build_input(self, mesh_file):
        """Convert a mesh file into tokenized model inputs, reducing size if needed.

        The function converts the mesh to a text representation, then
        tokenizes it and reduces the number of lines iteratively until the
        tokenized input fits within `max_input_tokens`.

        Args:
            mesh_file (str): Path to the mesh file to convert.

        Returns:
            dict: Tokenized tensors moved to the configured device.
        """
        mesh_string = self.mesh_renderer.mesh_to_string(mesh_file)
        reduced_lines = mesh_string.splitlines()
        original_line_count = len(reduced_lines)
        reduced_mesh = "\n".join(reduced_lines)
        messages = [
            {
                "role": "user",
                "content": f"What is this shape?\n{reduced_mesh}",
            },
            {
                "role": "system",
                "content": "A 3D model of "
            }
        ]
        inputs = self.tokenizer.apply_chat_template(messages, return_tensors="pt")

        while inputs["input_ids"].shape[-1] > self.max_input_tokens and len(reduced_lines) > 0:
            if len(reduced_lines) == 1:
                reduced_lines = []
            else:
                reduced_lines = reduced_lines[::2]
            reduced_mesh = "\n".join(reduced_lines)
            messages = [
                {
                    "role": "user",
                    "content": f"What is this shape?\n{reduced_mesh}",
                },
                {
                    "role": "system",
                    "content": "A 3D model of "
                }
            ]
            inputs = self.tokenizer.apply_chat_template(messages, return_tensors="pt")

        final_length = inputs["input_ids"].shape[-1]
        if len(reduced_lines) < original_line_count:
            print(
                f"Reduced mesh lines from {original_line_count} to {len(reduced_lines)} "
                f"to fit {final_length}/{self.max_input_tokens} input tokens."
            )

        device_inputs = inputs.to(self.device)

        del mesh_string, reduced_mesh, reduced_lines, messages, inputs
        gc.collect()

        return device_inputs

    def classify_one(self, file_path):
        """Classify a single mesh using the causal language model.

        Args:
            file_path (str): Path to the mesh file to classify.

        Returns:
            str: Lowercased predicted label.
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Mesh file {file_path} does not exist.")

        inputs = self.build_input(file_path)

        with torch.no_grad():
            outputs = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens, eos_token_id=self.terminators, do_sample=False, temperature=0.2, use_cache=True)

        input_length = inputs["input_ids"].shape[-1]
        outputs_trimmed = outputs[0][input_length:]
        predicted_label = self.tokenizer.decode(outputs_trimmed, skip_special_tokens=True).strip().lower()

        del inputs, outputs, outputs_trimmed
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
        gc.collect()

        print("\nMulti-view classification:")
        print(predicted_label)

        return predicted_label

    def classify_batch(self, folder_path, limit=None, save_path=None):
        """Classify meshes in a directory and optionally save results.

        Args:
            folder_path (str): Directory containing mesh files.
            limit (int, optional): Maximum number of files to process.
            save_path (str, optional): File path to write JSON results.

        Returns:
            str: Status message indicating completion and save location when provided.
        """
        if not os.path.isdir(folder_path):
            raise ValueError(f"Folder does not exist: {folder_path}")

        files = []
        for root, _, filenames in os.walk(folder_path):
            files.extend(os.path.join(root, f) for f in filenames)

        files.sort()

        if limit is not None:
            files = files[:limit]
        classifications = []
        label_processor = LabelProcessor()

        for file in files:
            print(f"\nClassifying {file}...")
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()
            predicted_label = self.classify_one(file)
            print(f"Predicted label: {predicted_label}")
            embedding = label_processor.compute_embedding(predicted_label)
            classifications.append(
                {
                    "file_path": file,
                    "predicted_label": predicted_label,
                    "embedding": embedding.tolist(),
                }
            )
            del embedding, predicted_label
            torch.cuda.empty_cache()

        if save_path:
            with open(save_path, "w") as f:
                json.dump(classifications, f)

        del label_processor
        torch.cuda.empty_cache()

        return (
            "Classification completed. Results saved to " + save_path
            if save_path
            else "Classification completed."
        )

    def classify(self, file_path=None, folder_path=None):
        """Dispatch method to classify a single file or a directory.

        Args:
            file_path (str, optional): Path to a single mesh file.
            folder_path (str, optional): Path to a directory of meshes.

        Returns:
            Any: Result returned by `classify_one` or `classify_batch`.

        Raises:
            ValueError: If neither argument is provided.
        """
        if file_path is not None:
            return self.classify_one(file_path)
        elif folder_path is not None:
            return self.classify_batch(folder_path)
        else:
            raise ValueError("Either file_path or folder_path must be provided.")