poc_model.py

# POC Model
from transformers.models.llama import LlamaForCausalLM

class POCModel(LlamaForCausalLM):
    def __init__(self, config):
        super().__init__(config)
        print("[POC] Model loaded")