FunctionGemma Domain Classifier

Fine-tuned FunctionGemma-270M for multi-domain query classification using LoRA.

Model Details

Base Model: google/functiongemma-270m-it
Model Size: 270M parameters (540MB)
Fine-tuning Method: LoRA (Low-Rank Adaptation)
Trainable Parameters: ~7.6M (2.75%)
Training Time: 23.3 minutes
Hardware: GPU (memory optimized for <5GB VRAM)

Performance

Accuracy: 95.51%
F1 Score (Weighted): 0.96
F1 Score (Macro): 0.88
Training Loss: 0.3

Supported Domains (17)

ambiguous
api_generation
business
coding
creative_content
data_analysis
education
general_knowledge
geography
history
law
literature
mathematics
medicine
science
sensitive
technology

Use Cases

Query Routing: Route user queries to specialized models/services
Content Classification: Categorize text by domain
Multi-domain Detection: Identify queries spanning multiple domains
Intent Analysis: Understand query context and domain

Quick Start

Installation

pip install transformers peft torch

Inference

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
import json

# Load model
base_model = AutoModelForCausalLM.from_pretrained(
    "google/functiongemma-270m-it",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, "ovinduG/functiongemma-domain-classifier")
tokenizer = AutoTokenizer.from_pretrained("ovinduG/functiongemma-domain-classifier")

# Classify a query
def classify(text):
    # Define function schema
    function_def = {
        "type": "function",
        "function": {
            "name": "classify_query_domain",
            "description": "Classify query into domains",
            "parameters": {
                "type": "object",
                "properties": {
                    "primary_domain": {"type": "string"},
                    "primary_confidence": {"type": "number"},
                    "is_multi_domain": {"type": "boolean"},
                    "secondary_domains": {"type": "array"}
                }
            }
        }
    }
    
    messages = [
        {"role": "developer", "content": "You are a model that can do function calling"},
        {"role": "user", "content": text}
    ]
    
    inputs = tokenizer.apply_chat_template(
        messages,
        tools=[function_def],
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt"
    ).to(model.device)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=False,
            pad_token_id=tokenizer.eos_token_id
        )
    
    response = tokenizer.decode(
        outputs[0][inputs["input_ids"].shape[-1]:],
        skip_special_tokens=True
    )
    
    # Parse function call
    if "{" in response:
        start = response.find("{")")
        end = response.rfind("}") + 1
        return json.loads(response[start:end])
    
    return {"error": "Failed to parse response"}

# Example
result = classify("Write a Python function to calculate fibonacci numbers")
print(json.dumps(result, indent=2))

Example Output

{
  "primary_domain": "coding",
  "primary_confidence": 0.95,
  "is_multi_domain": false,
  "secondary_domains": []
}

Multi-Domain Example

result = classify("Build an ML model to predict customer churn and create REST API endpoints")
print(json.dumps(result, indent=2))

{
  "primary_domain": "data_analysis",
  "primary_confidence": 0.85,
  "is_multi_domain": true,
  "secondary_domains": [
    {
      "domain": "api_generation",
      "confidence": 0.75
    }
  ]
}

Training Details

Dataset

Total Samples: 5,046
Training Samples: 3,666
Validation Samples: 690
Test Samples: 690
Multi-domain Queries: 546 (10.8%)

Training Configuration

# LoRA Configuration
r = 32
lora_alpha = 64
lora_dropout = 0.05
target_modules = ['q_proj', 'v_proj', 'k_proj', 'o_proj', 'gate_proj', 'up_proj', 'down_proj']

# Training Configuration  
num_epochs = 5
batch_size = 4
gradient_accumulation_steps = 8
learning_rate = 0.0003
max_length = 1024
optimizer = "adamw_8bit"  # Memory optimized

Memory Optimization

This model was trained with memory optimizations to run on GPUs with <5GB VRAM:

8-bit Optimizer: Reduces optimizer memory by 50%
Gradient Checkpointing: Trades compute for memory
Smaller Batches: 4 samples per batch with gradient accumulation
Shorter Sequences: 1024 tokens max (vs 2048)

Total VRAM Usage: ~4GB (vs ~40GB without optimization)

Performance by Domain

Domain	Precision	Recall	F1-Score	Support
ambiguous	0.98	1.00	0.99	45
api_generation	0.98	1.00	0.99	45
business	0.98	0.93	0.95	44
coding	0.98	0.96	0.97	48
creative_content	0.90	1.00	0.95	45
data_analysis	0.96	0.98	0.97	46
education	0.98	0.96	0.97	45
general_knowledge	0.76	0.84	0.80	45
law	0.98	0.94	0.96	49
literature	1.00	0.93	0.97	45
mathematics	1.00	1.00	1.00	47
medicine	0.98	0.89	0.93	46
science	1.00	0.98	0.99	47
sensitive	0.92	1.00	0.96	45
technology	1.00	0.93	0.97	46

Overall Accuracy: 95.51%

Advantages

✅ Tiny Size: 270M parameters (14x smaller than Phi-3)
✅ Fast Inference: 0.3s on CPU, 0.08s on GPU
✅ Low Memory: Runs on 4GB VRAM
✅ High Accuracy: 95.51% (competitive with larger models)
✅ Multi-domain: Detects queries spanning multiple domains
✅ Function Calling: Built-in structured output
✅ Mobile-Ready: Can deploy on smartphones

Limitations

Trained on English queries only
Performance varies by domain (see table above)
May struggle with highly ambiguous queries
Limited to 17 pre-defined domains

Base Model

Base model: google/functiongemma
Model family: Gemma
Model owner: Google LLC
Fine-tuning task: Domain classification

Acknowledgement & Attribution

This model is built upon Google’s FunctionGemma. Use of this model is subject to the Gemma Terms of Use and the Gemma Prohibited Use Policy:

Gemma Terms of Use (March 24, 2025):
https://ai.google.dev/gemma/terms
Gemma Prohibited Use Policy (February 21, 2024):
https://ai.google.dev/gemma/prohibited_use_policy

Users must comply with these policies when using, modifying, or distributing this model or its derivatives.

License

This model follows the same terms as Google’s Gemma models. Please review the above links for full license and usage restrictions.

Recommended Hugging Face Metadata

license: gemma
base_model: google/functiongemma
tags:
  - text-classification
  - domain-classification
  - gemma
  - functiongemma

Downloads last month: -; Downloads are not tracked for this model. How to track

Model tree for ovinduG/functiongemma-domain-classifier

Base model

google/functiongemma-270m-it

Adapter

(6)

this model

ovinduG
/

functiongemma-domain-classifier