Upload randtext.py

randtext.py

@@ -0,0 +1,126 @@
+import streamlit as st
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+st.set_page_config(
+    page_title="Section 8.1 Legal Assistant",
+    page_icon="⚖️",
+    layout="centered",
+)
+
+st.title("⚖️ Section 8.1 Legal Assistant")
+st.markdown("**Reinforcement Fine-Tuned Model for ITAA 1997 - Section 8.1 (General Deductions)**")
+st.markdown("---")
+
+SYSTEM_PROMPT = """You ONLY answer questions about Section 8.1 of the Income Tax Assessment Act 1997 (General Deductions). If a question is about any other section, topic, or contains wrong details about Section 8.1, refuse or correct it. Never add information not in Section 8.1."""
+
+MODEL_ID = "muhammadjasim12/rainforcejasim"
+
+
+@st.cache_resource
+def load_model():
+    """Load model once and cache it."""
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "Qwen/Qwen2.5-7B-Instruct",
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True,
+    )
+    model = PeftModel.from_pretrained(base_model, MODEL_ID)
+    model.eval()
+    return model, tokenizer
+
+
+def ask(question, model, tokenizer):
+    prompt = (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=300,
+            do_sample=False,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+    return tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[1]:],
+        skip_special_tokens=True
+    ).strip()
+
+
+# Load model
+with st.spinner("Loading model... (first time takes 2-3 minutes)"):
+    model, tokenizer = load_model()
+
+st.success("Model loaded!")
+st.markdown("---")
+
+# Chat interface
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat history
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"])
+
+# User input
+user_input = st.chat_input("Ask a question about Section 8.1...")
+
+if user_input:
+    # Show user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    with st.chat_message("user"):
+        st.markdown(user_input)
+
+    # Get model answer
+    with st.chat_message("assistant"):
+        with st.spinner("Thinking..."):
+            answer = ask(user_input, model, tokenizer)
+        st.markdown(answer)
+
+    st.session_state.messages.append({"role": "assistant", "content": answer})
+
+# Sidebar
+with st.sidebar:
+    st.header("About")
+    st.markdown("""
+    This model is **reinforcement fine-tuned (DPO + SFT)** 
+    exclusively on **Section 8.1** of the Income Tax 
+    Assessment Act 1997 (General Deductions).
+    
+    **It will:**
+    - Answer questions about Section 8.1
+    - Refuse questions about other sections
+    - Correct wrong details in questions
+    
+    **It will NOT:**
+    - Answer questions outside Section 8.1
+    - Add information not in the section
+    - Make up dollar amounts or rules
+    """)
+
+    st.markdown("---")
+    st.markdown("**Model:** `muhammadjasim12/rainforcejasim`")
+    st.markdown("**Base:** Qwen2.5-7B-Instruct")
+    st.markdown("**Training:** DPO + SFT")
+
+    st.markdown("---")
+    st.header("Example Questions")
+    st.markdown("""
+    - What is Section 8.1 about?
+    - What does Section 8.1(1)(a) say?
+    - Can I deduct a capital expense?
+    - What does Section 8.2 say?
+    - Does Section 8.1 have four subsections?
+    - What is Division 7A?
+    """)
+
+    if st.button("Clear Chat"):
+        st.session_state.messages = []
+        st.rerun()