import streamlit as st
import llama_cpp

@st.cache_resource
def load_model(model_path):
    return llama_cpp.Llama(model_path=model_path)

st.title("CyberSecurity Assistant")

model_path = "./llama-3-3b-cybersecurity-quantized.gguf"
temperature = st.sidebar.slider('Temperature', 0.0,2.0,0.2, step=0.1)

if model_path:
    try:
        llm = load_model(model_path)
        st.sidebar.success('Model loaded Successfully')
    except Exception as e:
        st.sidebar.error(f'Error loading the model: {e}')
        llm = None
else:
    st.warning("Model path not found")

# SYSTEM PROMPT

# GLOBAL VARIABLE INSTRUCTION
instruction= 'You are a Cybersecurity AI Assistant, will be glad to answer your questions related to Cybersecurity, particularly LLM Security.'


if llm:
    user_input = st.text_input("Your message", "")
    user_input= f'{instruction} \n\nUser: {user_input}\nAI'
    if user_input:
        with st.spinner("Generating response..."):
            try:
                response = llm(user_input, temperature=temperature, max_tokens=256)
                st.write(f"**Cybersecurity Assistant:** {response['choices'][0]['text'].strip()}")
            except Exception as e:
                st.error(f"Error in Generation: {e}")