Create app.py

app.py

@@ -0,0 +1,46 @@
+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/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.")