Update app.py

app.py

@@ -1,23 +1,57 @@
-import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+import streamlit as st
+
+# Load tokenizer and base model
+base_model_path = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(base_model_path)
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_path,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto" if torch.cuda.is_available() else None
+)
+
+# Load LoRA adapter
+model = PeftModel.from_pretrained(base_model, "lora_adapter")
+model.eval()
+
+# Streamlit UI
+st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
+st.write("Ask me any **Python programming** question:")
+
+user_input = st.text_input("Your question")
+
+if user_input:
+    # Prompt template that helps model decide to answer or reject
+    prompt = f"""
+You are a helpful and expert Python programming tutor.
+Answer only questions related to Python programming.
+If the question is unrelated to Python (like history, math, etc), politely respond:
+"Sorry, I can only answer Python-related questions."
+
+### Question:
+{user_input}
 
-st.title("TinyLLaMA Python Tutor (LoRA)")
+### Answer:
+"""
 
-@st.cache_resource
-def load_model():
-    base_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-    tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
-    model = PeftModel.from_pretrained(base_model, "lora_adapter")
-    return tokenizer, model
+    inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=True).to(model.device)
+    with torch.no_grad():
+        output = model.generate(
+            **inputs,
+            max_new_tokens=200,
+            temperature=0.7,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
-tokenizer, model = load_model()
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
 
-user_input = st.text_area("Ask me a Python coding question:")
+    # Clean the output to only show the answer
+    if "### Answer:" in decoded:
+        final_answer = decoded.split("### Answer:")[-1].strip()
+    else:
+        final_answer = decoded.strip()
 
-if st.button("Generate Answer"):
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=150)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    st.write("**Answer:**", response)
+    st.markdown(f"**Answer:** {final_answer}")