|
|
import streamlit as st |
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
import torch |
|
|
|
|
|
MODEL_NAME = "segolilylabs/Lily-Cybersecurity-7B-v0.2" |
|
|
|
|
|
@st.cache_resource() |
|
|
def load_model(): |
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) |
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto") |
|
|
return tokenizer, model |
|
|
|
|
|
tokenizer, model = load_model() |
|
|
|
|
|
st.title("🔐 Cybersecurity Chatbot") |
|
|
|
|
|
user_input = st.text_input("Ask a cybersecurity question:") |
|
|
|
|
|
if user_input: |
|
|
input_ids = tokenizer.encode(user_input, return_tensors="pt").to("cuda") |
|
|
output = model.generate(input_ids, max_length=200) |
|
|
response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True) |
|
|
|
|
|
st.write("🤖", response) |
|
|
|
