Code/inference_gemma3.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Gemma 3 12B - Fast Inference for Classification

Load the fine-tuned Gemma 3 model and run inference on test set.
Uses batch processing for faster inference.

Usage:
    python inference_gemma3.py
"""

import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"

import re
import torch
import pandas as pd
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel

# ---------------------------
# Paths & Config 
# ---------------------------
TEST_FILE = "/home/houssam-nojoom/.cache/huggingface/hub/datasets--houssamboukhalfa--telecom-ch1/snapshots/be06acac69aa411636dbe0e3bef5f0072e670765/test_file.csv"
BASE_MODEL = "google/gemma-3-4b-it"  # Must match training base model
ADAPTER_PATH = "./gemma3_classification_ft"  # LoRA adapter path

MAX_LENGTH = 2048
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Device: {DEVICE}")

# Enable TF32 for A100
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

text_col = "Commentaire client"

# ===========================================================================
# System Prompt and Few-Shot Examples (same as training)
# ===========================================================================
SYSTEM_PROMPT = """You are an expert Algerian linguist and data labeler. Your task is to classify customer comments from Algérie Télécom's social media into one of 9 specific categories.

## CLASSES (DETAILED DESCRIPTIONS)
- **Class 1 (Wish/Positive Anticipation):** Comments expressing a wish, a hopeful anticipation, or general positive feedback/appreciation for future services or offers.
- **Class 2 (Complaint: Equipment/Supply):** Comments complaining about the lack, unavailability, or delay in the supply of necessary equipment (e.g., modems, fiber optics devices).
- **Class 3 (Complaint: Marketing/Advertising):** Comments criticizing advertisements, marketing campaigns, or their lack of realism/meaning.
- **Class 4 (Complaint: Installation/Deployment):** Comments about delays, stoppages, or failure in service installation, network expansion, or fiber optics deployment (e.g., digging issues).
- **Class 5 (Inquiry/Request for Information):** Comments asking for eligibility, connection dates, service status, coverage details, or specific contact information.
- **Class 6 (Complaint: Technical Support/Intervention):** Comments regarding delays in repair interventions, issues with technical staff competence, or unsatisfactory customer service agency visits.
- **Class 7 (Pricing/Service Enhancement):** Comments focused on pricing, requests for cost reduction, or suggestions for general service/app functionality enhancements.
- **Class 8 (Complaint: Total Service Outage/Disconnection):** Comments indicating a complete, sustained loss of service (e.g., no phone, no internet, total disconnection).
- **Class 9 (Complaint: Service Performance/Quality):** Comments about technical issues impacting performance (e.g., slow speed, high latency, broken website/portal, coverage claims).

Respond with ONLY the class number (1-9). Do not include any explanation."""

FEW_SHOT_STRING = """
Comment: إن شاء الله يكون عرض صحاب 300 و 500 ميجا فيبر ياربي
Class: 1

Comment: الف مبروووك.. 
Class: 1

Comment:  - إتصالات الجزائر شكرا اتمنى لكم دوام الصحة والعافية 
Class: 1

Comment: C une fierté de faire partie de cette grande entreprise Algérienne de haute technologie et haute qualité
Class: 1

Comment: اتمنى لكم مزيد من التألق
Class: 1

Comment: زعما جابو المودام ؟
Class: 2

Comment: وفرو أجهزة مودام الباقي ساهل !
Class: 2

Comment: واش الفايدة تع العرض هذا هو اصلا لي مودام مهوش متوفر رنا قريب عام وحنا ستناو في جد موام هذا
Class: 2

Comment: Depuis un an et demi qu'on a installé w ma kan walou
Class: 2

Comment: قتلتونا بلكذب المودام غير متوفر عندي 4 أشهر ملي حطيت الطلب في ولاية خنشلة و مزال ماجابوش المودام
Class: 2

Comment: عندكم احساس و لا شريوه كما قالو خوتنا لمصريين
Class: 3

Comment: Kamel Dahmane الفايبر؟ مستحيل كامل عاجبتهم
Class: 3

Comment: ههههه نخلص مليون عادي كون يركبونا الفيبر 😂😂😂😂😂 كرهنا من 144p
Class: 3

Comment: إشهار بدون معنه
Class: 3

Comment: المشروع متوقف منذ اشهر
Class: 4

Comment: نتمنى تكملو في ايسطو وهران في اقرب وقت رانا نعانو مع ADSL
Class: 4

Comment: Fibre كاش واحد وصلوله الفيبر؟
Class: 4

Comment: ما هو الجديد وانا مزال ماعنديش الفيبر رغم الطلب ولالحاح
Class: 4

Comment: علبة الفيبر راكبة في الحي و لكن لا يوجد توصيل للمنزل للان
Class: 4

Comment: modem
Class: 5

Comment: يعني كي نطلعها ثلاثون ميغا كارطة تاع مائة الف قداه تحكملي؟
Class: 5

Comment: سآل الأماكن لي ما فيهاش الألياف البصرية إذا جابولنا الألياف السرعة تكون محدودة كيما ف ADSL؟
Class: 5

Comment: ماعرف كاش خبر على ايدوم 4G ماعرف تبقى قرد العش
Class: 5

Comment: هل متوفرة في حي عدل 1046 مسكن دويرة
Class: 5

Comment: عرض 20 ميجا نحيوه مدام مش قادرين تعطيونا حقنا
Class: 6

Comment: 4 سنوات وحنا نخلصو فالدار ماشفنا حتى bonus
Class: 6

Comment: لماذا التغيير في الرقم بدون تغيير سرعة التدفق هل من أجل الإشهار وفقط انا غير من 50 ميغا إلا 200 ميغا نظريا تغيرت وفي الواقع بقت قياس أقل من 50 ميغا
Class: 6

Comment: انا طلعت تدفق انترنات من 15 الى 20 عبر تطبيق my idoom لاكن سرعة لم تتغير
Class: 6

Comment: نقصوا الاسعار بزااااف غالية
Class: 7

Comment: علاه ماديروش في التطبيق خاصية التوقيف المؤقت للانترانات
Class: 7

Comment: وفرونا من بعد اي ساهلة
Class: 7

Comment: لازم ترجعو اتصال بتطبيقات الدفع بلا انترنت و مجاني ريقلوها يا اتصالات الجزائر
Class: 7

Comment: Promotion fin d'année ADSL idoom
Class: 7

Comment: رانا بلا تلفون ولا انترنت
Class: 8

Comment: ثلاثة اشهر بلا انترنت
Class: 8

Comment: votre site espace client ne fonctionne pas pourquoi?
Class: 8

Comment: ما عندنا الانترنيت ما نخلصوها من الدار
Class: 8

Comment: مشكل في 1.200جيق فيبر مدام نوكيا مخرج الانترنت 1جيق فقط كفاش راح تحلو هذا مشكل ومشكل ثاني فضاء الزبون ميمشيش مندو شهر
Class: 8

Comment: فضاء الزبون علاه منقدروش نسجلو فيه
Class: 9

Comment: هل موقع فضاء الزبون متوقف
Class: 9

Comment: ماراهيش توصل الفاتورة لا عن طريق الإيميل ولا عن طريق فضاء الزبون 
Class: 9

Comment: فضاء الزبون قرابة 20 يوم متوقف!!!!!!؟؟؟؟؟
Class: 9

Comment: برج الكيفان اظنها من العاصمة خارج تغطيتكم....احشموا بركاو بلا كذب....طلعنا الصواريخ للفضاء....بصح بالكذب....
Class: 9"""

# ===========================================================================
# Text Preprocessing
# ===========================================================================
def preprocess_text(text):
    """Preprocess text: remove tatweel, emojis, URLs, phone numbers."""
    if not isinstance(text, str):
        return ""
    
    # Remove URLs
    text = re.sub(r'https?://\S+|www\.\S+', '', text)
    
    # Remove email addresses
    text = re.sub(r'\S+@\S+', '', text)
    
    # Remove phone numbers
    text = re.sub(r'[\+]?[(]?[0-9]{1,4}[)]?[-\s\./0-9]{6,}', '', text)
    text = re.sub(r'\b0[567]\d{8}\b', '', text)
    text = re.sub(r'\b0[23]\d{7,8}\b', '', text)
    
    # Remove mentions
    text = re.sub(r'@\w+', '', text)
    
    # Remove Arabic tatweel
    text = re.sub(r'ـ+', '', text)
    
    # Remove emojis
    emoji_pattern = re.compile("["
        u"\U0001F600-\U0001F64F"
        u"\U0001F300-\U0001F5FF"
        u"\U0001F680-\U0001F6FF"
        u"\U0001F1E0-\U0001F1FF"
        u"\U00002702-\U000027B0"
        u"\U000024C2-\U0001F251"
        u"\U0001f926-\U0001f937"
        u"\U00010000-\U0010ffff"
        u"\u2640-\u2642"
        u"\u2600-\u2B55"
        u"\u200d"
        u"\u23cf"
        u"\u23e9"
        u"\u231a"
        u"\ufe0f"
        u"\u3030"
        "]+", flags=re.UNICODE)
    text = emoji_pattern.sub('', text)
    
    # Remove platform names
    text = re.sub(r'Algérie Télécom - إتصالات الجزائر', '', text, flags=re.IGNORECASE)
    text = re.sub(r'Algérie Télécom', '', text, flags=re.IGNORECASE)
    text = re.sub(r'إتصالات الجزائر', '', text)
    
    # Remove repeated characters
    text = re.sub(r'(.)\1{3,}', r'\1\1\1', text)
    
    # Normalize whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    
    return text

def format_prompt(comment):
    """Format prompt for inference."""
    user_prompt = f"""Here are some examples of how to classify comments:

{FEW_SHOT_STRING}

Now classify this comment:
Comment: {comment}
Class:"""
    return user_prompt

def create_inference_prompt(comment, tokenizer):
    """Create full prompt for inference."""
    clean_comment = preprocess_text(comment)
    
    messages = [
        {"role": "user", "content": SYSTEM_PROMPT + "\n\n" + format_prompt(clean_comment)}
    ]
    
    # Apply chat template with generation prompt
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    
    return text

def extract_class(generated_text):
    """Extract class number from generated text."""
    try:
        match = re.search(r'\b([1-9])\b', generated_text)
        if match:
            return int(match.group(1))
        return 1  # Default
    except:
        return 1

# ===========================================================================
# Main Inference
# ===========================================================================
print("\n" + "="*70)
print("Gemma 3 4B - Fast Batch Inference")
print("="*70 + "\n")

# Load tokenizer
print(f"Loading tokenizer from: {BASE_MODEL}")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id

tokenizer.padding_side = "left"  # Left padding for batch generation

# Load base model
print(f"Loading base model from: {BASE_MODEL}")
model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
    attn_implementation="eager",  # Flash attention not available
)

# Load LoRA adapter
print(f"Loading LoRA adapter from: {ADAPTER_PATH}")
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
model.eval()

print(f"\nModel loaded successfully!")

# Load test data
print(f"\nLoading test data from: {TEST_FILE}")
test_df = pd.read_csv(TEST_FILE)
print(f"Test samples: {len(test_df)}")

# Batch inference settings - use smaller batch for long prompts
BATCH_SIZE = 8  # Small batch to avoid OOM with long few-shot prompts

# Prepare all prompts first
print("\nPreparing prompts...")
all_prompts = []
for i in tqdm(range(len(test_df)), desc="Preparing"):
    comment = str(test_df.iloc[i][text_col])
    prompt = create_inference_prompt(comment, tokenizer)
    all_prompts.append(prompt)

# Run batch inference
print(f"\nRunning batch inference (batch_size={BATCH_SIZE})...")
all_preds = []
num_batches = (len(all_prompts) + BATCH_SIZE - 1) // BATCH_SIZE

with torch.no_grad():
    torch.cuda.empty_cache()  # Clear cache before inference
    
    for batch_idx in tqdm(range(num_batches), desc="Predicting"):
        start_idx = batch_idx * BATCH_SIZE
        end_idx = min(start_idx + BATCH_SIZE, len(all_prompts))
        batch_prompts = all_prompts[start_idx:end_idx]
        
        # Tokenize batch with padding
        inputs = tokenizer(
            batch_prompts,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_LENGTH,
            padding=True,  # Pad to longest in batch
        )
        input_lengths = [len(ids) for ids in inputs["input_ids"]]
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        # Generate batch (greedy decoding)
        outputs = model.generate(
            **inputs,
            max_new_tokens=3,  # Just need 1 digit
            do_sample=False,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
        
        # Decode each sequence in the batch
        for seq_idx, output_ids in enumerate(outputs):
            # Get only the new tokens (after the input)
            input_len = inputs["input_ids"].shape[1]  # All padded to same length
            generated_tokens = output_ids[input_len:]
            generated_text = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
            
            # Extract class
            pred_class = extract_class(generated_text)
            all_preds.append(pred_class)
        
        # Clear cache periodically
        if batch_idx % 50 == 0:
            torch.cuda.empty_cache()

# Save predictions
test_df["Predicted_Class"] = all_preds
output_file = "test_predictions_gemma3.csv"
test_df.to_csv(output_file, index=False)
print(f"\nPredictions saved to: {output_file}")

# Show sample predictions
print("\nSample predictions:")
for i in range(min(10, len(test_df))):
    text = str(test_df.iloc[i][text_col])
    text_display = text[:60] + "..." if len(text) > 60 else text
    pred = test_df.iloc[i]["Predicted_Class"]
    print(f"  [{i+1}] Class {pred}: {text_display}")

# Class distribution
print("\nPrediction distribution:")
pred_counts = test_df["Predicted_Class"].value_counts().sort_index()
for class_label, count in pred_counts.items():
    pct = count / len(test_df) * 100
    print(f"  Class {class_label}: {count:>5} samples ({pct:>5.1f}%)")

print("\n" + "="*70)
print("INFERENCE COMPLETE!")
print("="*70)
print(f"\nAdapter path: {ADAPTER_PATH}")
print(f"Test samples: {len(test_df)}")
print(f"Output file:  {output_file}")