|
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") |
|
|
|
|
|
|
|
|
|
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}") |
|
|