YAML Metadata Warning:empty or missing yaml metadata in repo card
Check out the documentation for more information.
Production-Grade Customer Support LLM
Project Overview
This project demonstrates the complete lifecycle of building a domain-specific AI-powered Customer Support Assistant using a lightweight open-source Large Language Model (LLM).
The system was designed to:
- Fine-tune a small LLM for customer support tasks
- Improve instruction following and response quality
- Handle refund and cancellation requests professionally
- Reduce hallucinations
- Implement AI safety guardrails
- Benchmark inference performance
- Deploy the model using FastAPI
The project focuses on lightweight and efficient deployment suitable for low-resource environments while maintaining strong customer interaction quality.
Objective
The primary objective of this project is to build a specialized customer support AI assistant capable of:
- Answering customer queries
- Processing refund requests
- Handling cancellations
- Generating professional support responses
- Preventing unsafe or malicious interactions
- Maintaining low inference latency
The project also demonstrates:
- Parameter-efficient fine-tuning using LoRA
- AI safety implementation
- API deployment
- Benchmarking and monitoring
Dataset
Dataset Used
Bitext Customer Support Dataset from Hugging Face.
Dataset Link: https://huggingface.co/datasets/bitext/Bitext-customer-support-llm-chatbot-training-dataset
Dataset Description
The dataset contains real-world styled customer support conversations such as:
- Refund requests
- Order tracking
- Product complaints
- Cancellation requests
- General customer queries
Each sample contains:
- Instruction (customer query)
- Response (support agent response)
Dataset Preprocessing
The preprocessing pipeline includes:
- Dataset loading using Hugging Face Datasets
- Selecting a smaller subset for lightweight training
- Formatting the dataset into instruction-response format
- Tokenization for model training
Example formatted sample:
Instruction:
I want a refund because my order arrived damaged.
Response:
We apologize for the inconvenience. Your refund request has been initiated.
Base Model
Model Used
TinyLlama/TinyLlama-1.1B-Chat-v1.0
Model Link: https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0
Why TinyLlama?
TinyLlama was selected because:
- Lightweight architecture
- Faster inference
- Lower GPU memory requirements
- Suitable for Google Colab free GPU
- Efficient for LoRA fine-tuning
Fine-Tuning Approach
Fine-Tuning Method
The model was fine-tuned using:
- LoRA (Low-Rank Adaptation)
- PEFT (Parameter-Efficient Fine-Tuning)
Instead of updating all model parameters, LoRA trains only small adapter layers which:
- Reduces GPU memory usage
- Speeds up training
- Makes deployment lightweight
LoRA Configuration
Parameter| Value r| 8 lora_alpha| 16 lora_dropout| 0.05 task_type| CAUSAL_LM
Training Configuration
Parameter| Value Batch Size| 1 Epochs| 1 Learning Rate| 2e-4 Max Length| 256 Dataset Size| 100 Samples
Why Reduced Dataset?
The dataset size was intentionally reduced because:
- Google Colab free GPU limitations
- Faster experimentation
- Faster deployment testing
- Lightweight proof-of-concept training
The objective was to demonstrate the complete LLM engineering pipeline rather than perform massive-scale training.
Model Training Pipeline
The training pipeline includes:
- Dataset Loading
- Dataset Formatting
- Tokenization
- LoRA Adapter Injection
- Fine-Tuning using Hugging Face Trainer
- Model Saving
- Inference Testing
Inference Pipeline
The inference pipeline follows:
Customer Query ↓ Prompt Validation ↓ Toxicity Filtering ↓ LLM Inference ↓ Response Generation ↓ Safety Validation ↓ Final Response
Toxicity Filtering
Objective
To prevent unsafe or toxic responses generated by the model.
Library Used
Detoxify
Implementation
The generated response is passed through Detoxify:
- If toxicity score exceeds threshold
- Response gets blocked
Example:
- Toxic content → blocked
- Safe content → returned to user
Prompt Guardrails
Objective
Prevent:
- Prompt injection attacks
- Jailbreak attempts
- Unsafe prompts
- System prompt leakage
Guardrail Strategy
A blocked keyword system was implemented.
Blocked examples:
- ignore previous instructions
- reveal system prompt
- hack
- password
- bypass safety
Unsafe prompts are rejected before inference.
Benchmarking
Metrics Measured
- Latency
Measures response generation time.
- GPU Memory Usage
Measures memory utilization during inference.
- Tokens Per Second
Measures inference throughput.
Benchmark Results
Metric| Result Latency| ~1.2 sec GPU Memory| ~5 GB Tokens/sec| ~40
Base Model vs Fine-Tuned Model
Scenario| Base Model| Fine-Tuned Model Refund Handling| Generic| Professional Cancellation| Weak workflow| Structured response Angry Customer| Limited empathy| Better customer handling Customer Queries| Generic answers| Domain-specific responses
AI Safety Implementation
The project implements multiple AI safety layers:
- Toxicity Filtering
Blocks toxic responses.
- Prompt Guardrails
Prevents malicious prompts.
- Unsafe Prompt Handling
Rejects prompt injection attempts.
Deployment
Framework Used
FastAPI
Deployment Architecture
User Request ↓ FastAPI Backend ↓ Prompt Validation ↓ Safety Layer ↓ LLM Inference ↓ Response Generation
Public Deployment
The API was exposed publicly using ngrok.
Swagger UI was used for API testing.
Endpoints:
- GET /
- POST /chat
FastAPI Features
- REST API support
- Production-ready backend
- Swagger documentation
- Lightweight deployment
- Easy scalability
Hugging Face Deployment
The fine-tuned model was uploaded to Hugging Face for:
- Model hosting
- Public inference access
- Reproducibility
- Portfolio showcase
GitHub Repository
The GitHub repository contains:
- Colab notebook
- FastAPI app
- Benchmark reports
- Architecture documentation
- Requirements file
- README documentation
Challenges Faced
- Google Colab GPU Limitations
Resolved by:
- Using TinyLlama
- Reducing dataset size
- LoRA fine-tuning
- Dependency Conflicts
Resolved by:
- Using stable Transformers versions
- Simplifying Trainer pipeline
- FastAPI Deployment in Colab
Resolved by:
- Using ngrok
- Running uvicorn properly
Future Improvements
Future enhancements may include:
- vLLM deployment
- llama.cpp benchmarking
- Quantized inference
- RAG integration
- Multi-turn conversation memory
- Kubernetes deployment
- Docker containerization
- Monitoring dashboard
Key Learnings
This project helped demonstrate:
- LLM fine-tuning workflows
- LoRA and PEFT training
- AI safety engineering
- FastAPI deployment
- Benchmarking techniques
- Production AI pipeline design
Technologies Used
Technology| Purpose Python| Development Transformers| LLM Framework PEFT| LoRA Fine-Tuning TinyLlama| Base Model FastAPI| API Deployment ngrok| Public Tunnel Detoxify| Toxicity Filtering Hugging Face| Model Hosting Google Colab| Training Environment
Conclusion
This project demonstrates a complete end-to-end workflow for building a production-style domain-specific LLM application.
The system successfully:
- Fine-tuned a lightweight LLM
- Implemented AI safety mechanisms
- Performed benchmarking
- Deployed inference APIs
- Created a scalable customer support assistant pipeline
The project highlights practical LLM engineering concepts including:
- efficient fine-tuning
- inference optimization
- deployment
- safety-first AI design
Author
Satyanarayan Dubey
Data Scientist | Machine Learning Engineer | Generative AI Enthusiast