| import gradio as gr |
| import spaces |
| import torch |
| from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
| |
| 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.""" |
|
|
|
|
| |
| 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 = """ |
| <div class="terminal-header"> |
| <div class="terminal-bar"> |
| <span class="terminal-dot dot-red"></span> |
| <span class="terminal-dot dot-yellow"></span> |
| <span class="terminal-dot dot-green"></span> |
| <span class="terminal-title">local-tech-advisor v1.0 — bash</span> |
| </div> |
| <div class="terminal-body"> |
| <span class="prompt">$</span> <span class="cmd">./local-tech-advisor --start</span><br> |
| <span class="output">[INFO] Loading model: <span class="cyan">nemotron-nano-4b</span> (fine-tuned)</span><br> |
| <span class="output">[INFO] Training cost: <span class="highlight">$2.03</span> | Training time: <span class="highlight">35 min</span></span><br> |
| <span class="output">[INFO] API calls needed: <span class="yellow">none. ever.</span></span><br> |
| <span class="output">[INFO] Status: <span class="cyan">ready</span> — ask me anything about AWS DevOps Agent</span><br> |
| <span class="prompt">$</span> <span class="cmd blink">_</span> |
| <div class="stats-row"> |
| <span class="stat-chip"><span class="val">4B</span> params</span> |
| <span class="stat-chip"><span class="val">66</span> docs</span> |
| <span class="stat-chip"><span class="val">1,230</span> Q&A pairs</span> |
| <span class="stat-chip"><span class="val">$0</span> per query</span> |
| <span class="stat-chip"><span class="val">0</span> APIs called</span> |
| </div> |
| </div> |
| </div> |
| """ |
|
|
| FOOTER = """ |
| <div class="footer-bar"> |
| <span class="prompt">$</span> echo "Built for Build Small Hackathon" | |
| model: <a href="https://huggingface.co/aslanconfig/tech-advisor-nemotron-4b" target="_blank">aslanconfig/tech-advisor-nemotron-4b</a> | |
| sponsor: nvidia | track: backyard |
| </div> |
| """ |
|
|
| |
| 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() |
|
|
