Upload Main.py

main.py

@@ -0,0 +1,58 @@
+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  # No match found
+
+# Function to generate response using the dataset or fallback to the pipeline
+def generate_response(user_input):
+    # First, try to find a match in the dataset
+    dataset_response = search_dataset(user_input)
+    
+    if dataset_response:
+        return dataset_response
+    
+    # If no match, fallback to the LLM
+    messages = [
+        {"role": "system", "content": "You are a pirate chatbot who specializes in Cisco switch and router configurations"},
+        {"role": "user", "content": user_input},
+    ]
+    
+    # Generate the response from the LLM
+    outputs = pipeline(messages, max_new_tokens=256)
+    
+    # Return the generated text
+    return outputs[0]["generated_text"]
+
+# Create Gradio interface
+iface = gr.Interface(
+    fn=generate_response,  # Function to call
+    inputs=gr.Textbox(lines=2, placeholder="Enter your Cisco switch/router question here..."),  # Updated Textbox input
+    outputs="text",  # Text output
+    title="Cisco Configuration Assistant",  # Title for the UI
+    description="Ask the chatbot questions about Cisco switch/router configurations",  # Description
+)
+
+# Launch the Gradio app
+iface.launch()