File size: 3,198 Bytes
4205fce 37b4b2f 4205fce 37b4b2f 4205fce 37b4b2f 4205fce af0aef7 4205fce | 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 | import transformers
import torch
import gradio as gr
from datasets import load_dataset
# Load the model once when the script starts
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
# Load the model into memory (on GPU if available)
pipeline = transformers.pipeline(
"text-generation",
model=model_id,
model_kwargs={"torch_dtype": torch.bfloat16},
device_map="auto", # Auto-detect GPU
)
# Load the dataset from Hugging Face
dataset = load_dataset("quantumminds/cisco_cli_commands")
# Function to search the dataset for a matching command
def search_dataset(user_input):
# Check if any command in the dataset matches the user input
for entry in dataset['train']: # assuming the dataset is in the 'train' split
if entry["command"].lower() in user_input.lower(): # Match the command with user input (case-insensitive)
return f"**Command:** {entry['command']}\n\n**Description:** {entry['description']}\n\n**Example:** {entry['examples'][0]['example_command'] if 'examples' in entry else 'No example available'}"
return None # If no match found
# Function to generate response using the dataset or fallback to the pipeline
def generate_response(user_input, chat_history):
# First, try to find a match in the dataset
dataset_response = search_dataset(user_input)
if dataset_response:
# Add user and assistant responses to the chat history from dataset match
chat_history.append({"role": "user", "content": user_input})
chat_history.append({"role": "assistant", "content": dataset_response})
return chat_history # Return early to avoid generating a response from the LLM
# If no match found in dataset, generate the response from the LLM
outputs = pipeline(user_input, max_new_tokens=512)
# Generate the assistant's response
assistant_response = outputs[0]["generated_text"]
# Add user and assistant responses to the chat history
chat_history.append({"role": "user", "content": user_input})
chat_history.append({"role": "assistant", "content": assistant_response})
return chat_history
# Create Gradio interface with chatbot and textbox
with gr.Blocks(theme=gr.themes.Ocean()) as iface:
gr.Markdown("<h1 style='text-align: center;'>Cisco Configuration Assistant</h1>")
chatbot = gr.Chatbot(label="Cisco Configuration Chatbot", type="messages", height=500)
user_input = gr.Textbox(placeholder="Enter your Cisco switch/router question here...", label="Your Input")
with gr.Row():
submit_btn = gr.Button("Submit")
clear_btn = gr.Button("Clear Feed")
def user(query, history):
# Generate a response and update the history
history = generate_response(query, history)
return history, "" # Return updated history and clear the input box
# Submit user input and update the chat history
user_input.submit(user, [user_input, chatbot], [chatbot, user_input])
submit_btn.click(user, [user_input, chatbot], [chatbot, user_input])
clear_btn.click(lambda: [], None, chatbot, queue=False)
# Prints to console if dataset is loaded
print(dataset)
# Launch the Gradio app
iface.launch()
|