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()