Views
No views yet

google/gemma-3-4b-it model that has been fine-tuned using QLoRA for a comprehensive, multi-task customer service application. The model was trained on a synthetic dataset of fashion-related customer complaints to perform both causal language modeling (generating a structured JSON response) and several classification tasks simultaneously via specialized classification heads.

is_actionable: Determines if the complaint requires a direct action (boolean).complaint_category: Classifies the complaint into one of 11 categories (e.g., "Sizing Issue", "Damaged Item").decision_recommendation: Recommends a course of action from 11 options (e.g., "Full_Refund_With_Return").info_complete: Assesses if all necessary information is present to resolve the issue (boolean).tone: Classifies the required tone for a formal response (e.g., "Empathetic_Standard").refund_percentage: Suggests a specific refund percentage (0-100).sentiment: Detects the customer's sentiment (e.g., "negative", "very_negative").aggression: Detects the level of aggression in the customer's message.
GemmaComplaintResolver wrapper class from the training notebook to be used correctly.1import torch
2from transformers import AutoTokenizer, AutoConfig
3from peft import PeftModel
4from huggingface_hub import hf_hub_download
5import os
6
7# You must have the GemmaComplaintResolver class definition in your environment.
8# Assuming it's defined as it was in the training notebook...
9
10# --- Configuration ---
11repo_id = "ShovalBenjer/gemma-3-4b-fashion-multitask_A4000_v7"
12device = "cuda" if torch.cuda.is_available() else "cpu"
13
14# --- 1. Load Tokenizer and Model Config ---
15tokenizer = AutoTokenizer.from_pretrained(repo_id)
16config = AutoConfig.from_pretrained("google/gemma-3-4b-it", trust_remote_code=True)
17
18# Define the label structure the model was trained with
19num_labels_dict = {
20 "is_actionable": 2, "complaint_category": 11, "decision_recommendation": 11,
21 "info_complete": 2, "tone": 7, "refund_percentage": 13,
22 "sentiment": 6, "aggression": 5
23}
24
25# --- 2. Instantiate the Custom Model Wrapper ---
26# IMPORTANT: This assumes the GemmaComplaintResolver class is defined.
27model = GemmaComplaintResolver(
28 base_model_name_or_path="google/gemma-3-4b-it",
29 num_labels_dict=num_labels_dict,
30 model_config_for_base_loading=config,
31)
32
33# --- 3. Load the Fine-Tuned Weights ---
34# a) Load the classification head weights
35weights_path = hf_hub_download(repo_id=repo_id, filename="classification_heads.pth")
36model.load_state_dict(torch.load(weights_path, map_location='cpu'), strict=False)
37
38
39# b) Apply the LoRA adapter
40model = PeftModel.from_pretrained(model, repo_id)
41
42# --- 4. Prepare for Inference ---
43# Cast to appropriate dtype and move to device
44compute_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
45model.to(dtype=compute_dtype).to(device).eval()
46
47# --- 5. Run Inference ---
48customer_complaint = "The t-shirt I ordered arrived with a huge hole in it! I'm very angry and want a full refund immediately."
49# The model expects the full prompt structure used during training.
50# In this notebook, the pre-processed column was 'text_for_lm'.
51# The structure inside 'text_for_lm' was: <start_of_turn>user\n{complaint_details}<end_of_turn>\n<start_of_turn>model\n{json_output}<eos>
52# For inference on just the classification heads, we only need the prompt part.
53input_text = f"<start_of_turn>user\\n{customer_complaint}<end_of_turn>\\n<start_of_turn>model\\n"
54
55inputs = tokenizer(input_text, return_tensors="pt").to(device)
56
57with torch.no_grad():
58 outputs = model(**inputs)
59
60# --- 6. Decode a Prediction ---
61# Example: Get the predicted complaint category
62category_logits = outputs['logits_complaint_category']
63predicted_category_id = torch.argmax(category_logits, dim=-1).item()
64complaint_categories = ["Sizing Issue", "Damaged Item", "Not as Described", "Shipping Problem", "Policy Inquiry", "Late Delivery", "Wrong Item Received", "Quality Issue", "Return Process Issue", "Other", "N/A"]
65predicted_category = complaint_categories[predicted_category_id]
66
67print(f"Customer Complaint: '{customer_complaint}'")
68print(f"Predicted Complaint Category: {predicted_category}")