File size: 798 Bytes
42ae955
 
 
 
bfdb98d
42ae955
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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)