Update app.py

app.py

@@ -3,35 +3,25 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
 import streamlit as st
 
-# Load base model and tokenizer
+# Load tokenizer and model
 base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = AutoTokenizer.from_pretrained(base_model)
 
-# Load base model on CPU
-model = AutoModelForCausalLM.from_pretrained(
-    base_model,
-    torch_dtype=torch.float32,
-    device_map="cpu"
-)
-
-# Load LoRA adapter
-model = PeftModel.from_pretrained(model, "lora_adapter", device_map="cpu")
+model = AutoModelForCausalLM.from_pretrained(base_model, device_map="cpu")
+model = PeftModel.from_pretrained(model, "lora_adapter")  # Change if needed
 model.eval()
 
-# Prompt template
+# Format prompt (no USER/ASSISTANT lines to confuse model)
 def format_prompt(instruction):
-    return f"""### SYSTEM:
-You are a helpful and expert Python programming tutor.
+    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:
 "Sorry, I can only answer Python-related questions."
 
-### USER:
-{instruction}
-
-### ASSISTANT:
-"""
+Question: {instruction}
+Answer:"""
 
+# Generate answer
 def chat(instruction):
     prompt = format_prompt(instruction)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
@@ -48,17 +38,12 @@ def chat(instruction):
 
     full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only the final answer
-    if "### Answer:" in full_output:
-        return full_output.split("### Answer:")[-1].strip()
-    elif "### ASSISTANT:" in full_output:
-        return full_output.split("### ASSISTANT:")[-1].strip()
-    else:
-        return full_output.strip()
+    # Only return the model's answer
+    return full_output.split("Answer:")[-1].strip()
 
 # Streamlit UI
-st.set_page_config(page_title="🐍 Python Tutor Chatbot")
-st.title("🐍 Python Tutor Chatbot(LoRA")
+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:")
@@ -67,4 +52,7 @@ if st.button("Get Answer") and user_input.strip():
     with st.spinner("Thinking..."):
         response = chat(user_input)
         st.markdown("**Answer:**")
-        st.write(response)
+        if "```" in response:
+            st.markdown(response)
+        else:
+            st.write(response)