Update app.py

app.py

@@ -1,65 +1,58 @@
-import os
 import torch
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
-
-# Set CPU device
-device = torch.device("cpu")
+import streamlit as st
 
 # Load base model and tokenizer
-model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+base_model = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(base_model)
 
-base_model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float32,  # You can try float16 if supported
-    device_map={"": device}
+# Load base model in CPU-only mode
+model = AutoModelForCausalLM.from_pretrained(
+    base_model,
+    device_map="cpu"
 )
 
-# Load LoRA adapter (your fine-tuned weights)
-model = PeftModel.from_pretrained(base_model, "lora_adapter", device_map={"": device})
+# Load the LoRA adapter
+model = PeftModel.from_pretrained(model, "lora_adapter")
 model.eval()
 
-# Format the prompt with filtering
 def format_prompt(instruction):
-    return f"""You are a helpful and expert Python programming tutor.
-Only answer questions related to Python programming.
-If the question is unrelated to Python, respond with:
-"Sorry, I can only answer Python-related questions."
+    return f"""### SYSTEM:
+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.\"
 
-### Instruction:
+### USER:
 {instruction}
 
-### Response:
+### ASSISTANT:
 """
 
-# Generate answer
 def chat(instruction):
     prompt = format_prompt(instruction)
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
         **inputs,
         max_new_tokens=256,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
+        do_sample=False,
+        temperature=0.0,
+        top_p=1.0,
         repetition_penalty=1.1
     )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response.split("### Response:")[-1].strip()
+    return response.split("### ASSISTANT:")[-1].strip()
 
 # Streamlit UI
+st.set_page_config(page_title="🐍 Python Tutor Chatbot")
 st.title("🐍 Python Tutor Chatbot")
-st.markdown("Ask me Python programming questions!")
+st.write("Ask me Python programming questions!")
 
-question = st.text_area("Your question:")
+user_input = st.text_area("Your question:")
 
-if st.button("Answer"):
-    if question.strip():
-        with st.spinner("Thinking..."):
-            response = chat(question)
+if st.button("Get Answer") and user_input.strip():
+    with st.spinner("Thinking..."):
+        response = chat(user_input)
         st.markdown("**Answer:**")
         st.write(response)
-    else:
-        st.warning("Please enter a question.")