Update app.py

app.py

@@ -1,53 +1,55 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import streamlit as st
 
-st.set_page_config(page_title="TinyLLaMA Python Tutor", layout="centered")
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0")
 
-st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
-st.write("Ask me any Python programming question:")
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    "TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto"
+)
 
-@st.cache_resource
-def load_model():
-    base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    adapter_path = "lora_adapter"
+# Load LoRA adapter
+model = PeftModel.from_pretrained(base_model, "lora_adapter")
 
-    tokenizer = AutoTokenizer.from_pretrained(base_model)
-    model = AutoModelForCausalLM.from_pretrained(base_model, torch_dtype=torch.float32)
-    model = PeftModel.from_pretrained(model, adapter_path)
-    model.eval()
+# Set title
+st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
+st.markdown("Ask me any **Python programming** question:")
 
-    return tokenizer, model
+# User input
+user_question = st.text_input("Your question")
 
-tokenizer, model = load_model()
+if user_question:
+    with st.spinner("Thinking..."):
 
-def build_prompt(question):
-    return (
-        "You are a helpful and concise Python programming tutor. "
-        "If the question is not about Python, respond with: "
-        "'Sorry, I can only answer Python-related questions.'\n\n"
-        f"Question: {question}\nAnswer:"
-    )
+        # Clean prompt
+        prompt = f"""
+You are a helpful and expert Python programming tutor.
+If the question is about Python, explain clearly with examples.
+If the question is unrelated to Python, respond with "Sorry, I can only answer Python-related questions."
 
-question = st.text_input("Your question")
+Question: {user_question}
+Answer:"""
 
-if question:
-    prompt = build_prompt(question)
-    inputs = tokenizer(prompt, return_tensors="pt")
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-    with st.spinner("Thinking..."):
-        outputs = model.generate(
+        output = model.generate(
             **inputs,
-            max_new_tokens=250,
-            temperature=0.6,
-            top_p=0.85,
-            repetition_penalty=1.2,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
+            max_new_tokens=256,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.7,
+            repetition_penalty=1.1
         )
 
-    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = decoded_output.split("Answer:")[-1].strip()
+        decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
+
+        # Extract only the generated answer after "Answer:"
+        answer_start = decoded_output.find("Answer:")
+        answer = decoded_output[answer_start + len("Answer:"):].strip() if answer_start != -1 else decoded_output.strip()
 
-    st.markdown(f"**💬 Answer:**\n\n{answer}")
+        st.markdown(f"💬 **Answer:**\n\n{answer}")