import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
# --- Configuration ---
MODEL_ID = "aslanconfig/tech-advisor-nemotron-4b"
SYSTEM_PROMPT = """/no_think
You are Tech Advisor, an expert on AWS cloud services with deep knowledge of AWS DevOps Agent.
You have comprehensive knowledge of AWS DevOps Agent including:
- What it is and how it works (Agent Spaces, topology, dual-console architecture)
- Key features: autonomous incident response, proactive prevention, on-demand SRE tasks
- Integrations: CloudWatch, Datadog, Dynatrace, New Relic, Splunk, Grafana, PagerDuty, GitHub, GitLab, Azure DevOps, ServiceNow, Slack
- GA features: Azure/on-prem support, Triage Agent, Learned/Custom Skills, Code Indexing, Private Connections
- Pricing: $0.0083 per agent-second, free trial details, AWS Support credits
- Getting started: Agent Spaces, connecting tools, running investigations
- Security: encryption, customer managed keys, IdP integration, CloudTrail auditing
- Available regions: US East, US West, Frankfurt, Ireland, Sydney, Tokyo
Be concise and structured. Use bullet points where appropriate. Provide accurate, detailed answers."""
# --- Load model ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
).to(device).eval()
@spaces.GPU
def respond(message: str, history: list[dict]) -> str:
"""Generate a response using the fine-tuned Llama-3.1-Nemotron-Nano-4B."""
messages = [{"role": "system", "content": SYSTEM_PROMPT}]
for msg in history:
if msg.get("role") and msg.get("content"):
messages.append({"role": msg["role"], "content": str(msg["content"])})
messages.append({"role": "user", "content": message})
input_ids = tokenizer.apply_chat_template(
messages, return_tensors="pt", add_generation_prompt=True, tokenize=True
)
if hasattr(input_ids, "input_ids"):
input_ids = input_ids.input_ids
if not isinstance(input_ids, torch.Tensor):
input_ids = torch.tensor([input_ids])
input_ids = input_ids.to(model.device)
with torch.no_grad():
generated_ids = model.generate(
input_ids,
max_new_tokens=2048,
do_sample=True,
temperature=0.3,
top_p=0.9,
eos_token_id=tokenizer.eos_token_id,
)
output_text = tokenizer.decode(
generated_ids[0][input_ids.shape[1]:],
skip_special_tokens=True,
)
return output_text
CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;700&display=swap');
.gradio-container {
max-width: 900px !important;
margin: auto !important;
background: #0d1117 !important;
font-family: 'JetBrains Mono', monospace !important;
min-height: 100vh;
}
footer { display: none !important; }
.terminal-header {
background: #161b22;
border: 1px solid #30363d;
border-radius: 8px;
padding: 0;
margin-bottom: 20px;
overflow: hidden;
box-shadow: 0 8px 32px rgba(0, 255, 65, 0.1);
}
.terminal-bar {
background: #21262d;
padding: 8px 14px;
display: flex;
align-items: center;
gap: 8px;
border-bottom: 1px solid #30363d;
}
.terminal-dot {
width: 12px; height: 12px; border-radius: 50%;
}
.dot-red { background: #ff5f56; }
.dot-yellow { background: #ffbd2e; }
.dot-green { background: #27c93f; }
.terminal-title {
color: #8b949e;
font-size: 0.8em;
margin-left: 10px;
font-family: 'JetBrains Mono', monospace;
}
.terminal-body {
padding: 20px 24px;
color: #c9d1d9;
font-family: 'JetBrains Mono', monospace;
font-size: 0.9em;
line-height: 1.8;
}
.terminal-body .prompt { color: #27c93f; }
.terminal-body .cmd { color: #f0f6fc; font-weight: 700; }
.terminal-body .output { color: #8b949e; }
.terminal-body .highlight { color: #ff7b72; }
.terminal-body .cyan { color: #79c0ff; }
.terminal-body .yellow { color: #e3b341; }
.stats-row {
display: flex;
gap: 12px;
margin-top: 16px;
flex-wrap: wrap;
}
.stat-chip {
background: #21262d;
border: 1px solid #30363d;
border-radius: 6px;
padding: 6px 12px;
font-size: 0.8em;
color: #c9d1d9;
font-family: 'JetBrains Mono', monospace;
}
.stat-chip .val { color: #27c93f; font-weight: 700; }
.footer-bar {
text-align: center;
padding: 14px;
margin-top: 15px;
background: #161b22;
border: 1px solid #30363d;
border-radius: 8px;
color: #8b949e;
font-size: 0.8em;
font-family: 'JetBrains Mono', monospace;
}
.footer-bar a { color: #27c93f; text-decoration: none; }
.footer-bar a:hover { text-decoration: underline; }
"""
HEADER = """
"""
FOOTER = """
"""
# --- Gradio UI ---
with gr.Blocks(css=CSS, title="Local Tech Advisor") as demo:
gr.HTML(HEADER)
gr.ChatInterface(
fn=respond,
examples=[
"What is AWS DevOps Agent and how does it work?",
"How much does AWS DevOps Agent cost?",
"What observability tools does it integrate with?",
"How do I create an Agent Space?",
"What security features does AWS DevOps Agent provide?",
"What regions is AWS DevOps Agent available in?",
],
cache_examples=False,
)
gr.HTML(FOOTER)
if __name__ == "__main__":
demo.launch()