File size: 3,940 Bytes
9394f3f | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """
Evaluation script for the IaC-Coder model.
Tests the model on representative IaC tasks across Terraform, Ansible, Docker, K8s, and Shell.
Usage:
pip install transformers torch peft
python eval_iac.py
"""
import torch
from transformers import pipeline
MODEL_ID = "Tejas86/iac-coder-1.5b"
print(f"Loading model from {MODEL_ID}...")
pipe = pipeline(
"text-generation",
model=MODEL_ID,
torch_dtype=torch.bfloat16,
device_map="auto",
)
test_cases = [
{
"name": "Terraform: Create an S3 bucket with versioning",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Terraform HCL code."},
{"role": "user", "content": "Task: Create an AWS S3 bucket with versioning enabled and server-side encryption\n\nGenerate the Terraform HCL code."},
],
},
{
"name": "Terraform: Modify VPC to add private subnet",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Terraform HCL code."},
{"role": "user", "content": 'Task: Add a private subnet to the VPC\n\nModify the following Terraform HCL file:\n```\nresource "aws_vpc" "main" {\n cidr_block = "10.0.0.0/16"\n tags = {\n Name = "main-vpc"\n }\n}\n\nresource "aws_subnet" "public" {\n vpc_id = aws_vpc.main.id\n cidr_block = "10.0.1.0/24"\n tags = {\n Name = "public-subnet"\n }\n}\n```'},
],
},
{
"name": "Ansible: Install and configure Nginx",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready YAML (Ansible/K8s) code."},
{"role": "user", "content": "Task: Write an Ansible playbook to install nginx on Ubuntu and start the service\n\nGenerate the YAML (Ansible/K8s) code."},
],
},
{
"name": "Kubernetes: Redis StatefulSet with PVC",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready YAML (Ansible/K8s) code."},
{"role": "user", "content": "Task: Create a Kubernetes StatefulSet for Redis with persistent storage\n\nGenerate the YAML (Ansible/K8s) code."},
],
},
{
"name": "Dockerfile: Multi-stage Python Flask app",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Dockerfile code."},
{"role": "user", "content": "Task: Create a multi-stage Dockerfile for a Python Flask application with pip dependencies\n\nGenerate the Dockerfile code."},
],
},
{
"name": "Shell: CI/CD deployment script",
"messages": [
{"role": "system", "content": "You are an expert DevOps engineer. Generate clean, production-ready Shell/Bash script code."},
{"role": "user", "content": "Task: Write a bash script that builds a Docker image, runs tests, and pushes to ECR\n\nGenerate the Shell/Bash script code."},
],
},
]
print("\n" + "=" * 80)
print("IaC-CODER MODEL EVALUATION")
print("=" * 80)
for i, test in enumerate(test_cases, 1):
print(f"\n{'─' * 80}")
print(f"Test {i}/{len(test_cases)}: {test['name']}")
print(f"{'─' * 80}")
output = pipe(
test["messages"],
max_new_tokens=512,
temperature=0.2,
do_sample=True,
top_p=0.9,
)
generated = output[0]["generated_text"]
if isinstance(generated, list):
assistant_msgs = [m for m in generated if m["role"] == "assistant"]
response = assistant_msgs[-1]["content"] if assistant_msgs else str(generated)
else:
response = generated
print(f"\nGenerated:\n{response[:1500]}")
print(f"\n{'=' * 80}")
print(f"Evaluation complete! Model: {MODEL_ID}")
print(f"{'=' * 80}")
|