Spaces:
Sleeping
Sleeping
File size: 1,408 Bytes
46977f8 a5617d9 46977f8 |
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 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import streamlit as st
from unsloth import FastLanguageModel
from transformers import TextStreamer
# Page Configuration
st.set_page_config(page_title="AI Traffic Law Advisor", layout="wide")
# Load the LoRA model
MODEL_PATH = "./lora_model"
@st.cache_resource(show_spinner=False)
def load_model():
# Load model and tokenizer
model, tokenizer = FastLanguageModel.from_pretrained(
MODEL_PATH,
device_map="auto"
)
# Enable inference mode
model = FastLanguageModel.for_inference(model)
return model, tokenizer
model, tokenizer = load_model()
st.title("AI Traffic Law Advisor")
user_query = st.text_area("Enter your legal question about traffic rules in India:", "")
if st.button("Get Advice"):
if user_query.strip():
messages = [{"role": "user", "content": user_query}]
# Tokenize input
inputs = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_tensors="pt"
).to(model.device)
# Stream response
text_streamer = TextStreamer(tokenizer, skip_prompt=True)
st.markdown("**AI Response:**")
with st.spinner("Generating response..."):
model.generate(input_ids=inputs, streamer=text_streamer, max_new_tokens=1048, temperature=0.7)
else:
st.warning("Please enter a query.")
|