import torch
from transformers import BartForConditionalGeneration, BartTokenizer
import gradio as gr

 

# Load the fine-tuned model and tokenizer
model_path = "difinative/AIBuddy"  # Path to the pretrained fine-tuned model
model = BartForConditionalGeneration.from_pretrained(model_path)
tokenizer = BartTokenizer.from_pretrained(model_path)

 

# Translate function using the fine-tuned model
def translate_instruction(context, input_text):
    full_input = context + " " + input_text
    input_ids = tokenizer.encode(full_input, truncation=True, return_tensors='pt')
    with torch.no_grad():
        outputs = model.generate(input_ids, max_length=50, num_beams=4, early_stopping=True)
    translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_text
 

# Gradio Interface
iface = gr.Interface(
    fn=translate_instruction,
    inputs=[
        gr.inputs.Textbox(label="Enter Context"),
        gr.inputs.Textbox(label="Enter Human Instruction")
    ],
    outputs=gr.outputs.Textbox(label="Generated CLI Command"),
    title="CLI Command Generator",
    description="<p>This tool generates CLI commands from human instructions and context.</p>"
                "<p>Provide a context and human instruction, and click 'Generate' to get the CLI command.</p>",
    examples=[
        ["Pod xyz is called as tiger.", "Show me a list of pods in the current namespace."],
        ["Nickname of deployment development is cap", "Display the pods running in the cap."],
        ["Production is also known as", "List all pods in the 'production' environment."],
        ["Kubernetes", "Show the pods running in the 'testing' namespace?"],
        ["Docker", "What command should I use to get the list of containers?"]
    ]
)

 

# Launch the Gradio app using Ngrok
iface.launch()