Update app.py

app.py

@@ -1,46 +1,58 @@
+import torch
+from peft import AutoPeftModelForCausalLM
+from transformers import AutoTokenizer, TextStreamer
 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
+
+# Initialize Streamlit UI
+st.title("Legal Query Chatbot")
+st.write("Ask questions related to Indian traffic laws and get AI-generated responses.")
+
+# Load LoRA fine-tuned model and tokenizer
+model_path = "lora_model"
+load_in_4bit = True
+
+# Load the model
+model = AutoPeftModelForCausalLM.from_pretrained(
+    model_path,
+    torch_dtype=torch.float16 if not load_in_4bit else torch.float32,
+    load_in_4bit=load_in_4bit,
+    device_map="auto"
+)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# Enable inference mode
+model.eval()
+
+# Streamlit input for user prompt
+user_input = st.text_input("Enter your legal query:", "What are the penalties for breaking a red light in India?")
+
+if user_input:
+    # Prepare the prompt
+    messages = [{"role": "user", "content": user_input}]
+
+    # Tokenize input
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+        return_tensors="pt"
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Streamlit progress indicator
+    with st.spinner("Generating response..."):
+        # Use a text streamer for efficient streaming output
         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.")
+
+        # Generate response
+        output = model.generate(
+            input_ids=inputs,
+            streamer=text_streamer,
+            max_new_tokens=128,
+            use_cache=True,
+            temperature=1.5,
+            min_p=0.1
+        )
+
+    st.success("Generation Complete!")