Changed dataset to be pulled from hugging face

main.py

@@ -1,58 +1,60 @@
-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()
+import transformers
+import torch
+import gradio as gr
+from datasets import load_dataset
+
+# Remember to add access token to huggingface-cli login
+
+# 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):
+    # 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..."),
+    outputs="text",
+    title="Cisco Configuration Assistant",  # Title for the UI
+    description="Ask the chatbot questions about Cisco switch/router configurations",
+)
+
+# Launch the Gradio app
+iface.launch()