File size: 1,762 Bytes
5037564 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
class SentimentAnalysisHandler:
def __init__(self):
"""Load base model and fine-tuned adapter."""
self.base_model_id = "unsloth/llama-3-8b-bnb-4bit"
self.adapter_model_id = "samiur-r/BanglishSentiment-Llama3-8B"
# Load tokenizer
self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
# Load base model
self.model = AutoModelForCausalLM.from_pretrained(
self.base_model_id,
device_map="auto",
torch_dtype=torch.bfloat16
)
# Attach LoRA adapter
self.model = PeftModel.from_pretrained(self.model, self.adapter_model_id)
self.model.eval()
def preprocess(self, input_text):
"""Tokenize input text."""
inputs = self.tokenizer(input_text, return_tensors="pt").to("cuda")
return inputs
def inference(self, inputs):
"""Perform model inference."""
with torch.no_grad():
output = self.model.generate(**inputs, max_new_tokens=256)
return output
def postprocess(self, output):
"""Decode model output."""
sentiment = self.tokenizer.decode(output[0], skip_special_tokens=True)
return sentiment
def predict(self, input_text):
"""Full prediction pipeline."""
inputs = self.preprocess(input_text)
output = self.inference(inputs)
return self.postprocess(output)
# Create handler instance
_model_handler = SentimentAnalysisHandler()
def handle(inputs, context):
"""Entry point for model API inference."""
text = inputs.get("text", "")
return {"prediction": _model_handler.predict(text)}
|