Update app.py

app.py

@@ -1,58 +1,64 @@
+import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
-import streamlit as st
 
-# Load tokenizer and model
-base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(base_model)
+# Load base model & tokenizer
+base_model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+adapter_path = "lora_adapter"  # path to your LoRA adapter directory
+
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+    base_model = AutoModelForCausalLM.from_pretrained(base_model_name, device_map="auto")
+    model = PeftModel.from_pretrained(base_model, adapter_path)
+    model.eval()
+    return tokenizer, model
+
+tokenizer, model = load_model()
+
+# Prompt formatting
+def format_prompt(user_input):
+    return f"""You are a helpful and knowledgeable Python tutor chatbot.
 
-model = AutoModelForCausalLM.from_pretrained(base_model, device_map="cpu")
-model = PeftModel.from_pretrained(model, "lora_adapter")  # Change if needed
-model.eval()
+You only answer questions related to Python programming, including:
+- Python syntax, functions, loops, and conditionals
+- Standard libraries and popular packages (e.g., NumPy, pandas)
+- Debugging and code explanation
+- Python tools, environments, and tips
 
-# Format prompt (no USER/ASSISTANT lines to confuse model)
-def format_prompt(instruction):
-    return f"""You are a helpful and expert Python programming tutor.
-You only answer questions related to Python programming.
-If the question is unrelated to Python, say:
+If a question is not related to Python, reply with:
 "Sorry, I can only answer Python-related questions."
 
-Question: {instruction}
-Answer:"""
+### Instruction:
+{user_input}
 
-# Generate answer
-def chat(instruction):
-    prompt = format_prompt(instruction)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+### Response:"""
 
+# Chat handler
+def chat(user_input):
+    prompt = format_prompt(user_input)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     with torch.no_grad():
-        outputs = model.generate(
+        output = model.generate(
             **inputs,
-            max_new_tokens=512,
-            do_sample=False,
+            max_new_tokens=200,
+            do_sample=True,
             temperature=0.7,
             top_p=0.9,
-            repetition_penalty=1.1
+            pad_token_id=tokenizer.eos_token_id
         )
-
-    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    # Only return the model's answer
-    return full_output.split("Answer:")[-1].strip()
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    return decoded.split("### Response:")[-1].strip()
 
 # Streamlit UI
-st.set_page_config(page_title="🐍 Python Tutor Chatbot (LoRA)")
-st.title("🐍 Python Tutor Chatbot (LoRA)")
-st.write("Ask me Python programming questions!")
-
-user_input = st.text_area("Your question:")
-
-if st.button("Get Answer") and user_input.strip():
-    with st.spinner("Thinking..."):
-        response = chat(user_input)
-        st.markdown("**Answer:**")
-        if "```" in response:
-            st.markdown(response)
-        else:
-            st.write(response)
+st.title("🧑‍🏫 Python Tutor Chatbot")
+st.write("Ask me anything about Python programming!")
+
+user_input = st.text_area("Your Question", height=150)
+if st.button("Ask"):
+    if user_input.strip():
+        with st.spinner("Thinking..."):
+            answer = chat(user_input)
+        st.markdown("### 💡 Answer:")
+        st.write(answer)