Update app.py

app.py

@@ -1,48 +1,78 @@
-import streamlit as st
+import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import streamlit as st
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig
+)
+from peft import PeftModel
 
-# Load model and tokenizer
-model_name = "lora_adapter"  # Update this to your LoRA model path
+# Offload directory for CPU inference
+os.makedirs("offload", exist_ok=True)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
+# Load base model + quantization config
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 
-# Chat function with prompt-based filtering
-def chat(instruction):
-    prompt =  f"""You are a helpful and expert Python programming tutor.
+bnb_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0,
+    llm_int8_enable_fp32_cpu_offload=True
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+base_model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=bnb_config,
+    device_map="auto",
+    offload_folder="offload"
+)
+
+# Load LoRA adapter
+model = PeftModel.from_pretrained(base_model, "lora_adapter")
+
+# Evaluation mode
+model.eval()
+
+# Prompt template
+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."
 
-Answer the following instruction with a structured and informative format.
 ### Instruction:
 {instruction}
 
 ### Response:
 """
+
+# Chat function
+def chat(instruction):
+    prompt = format_prompt(instruction)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=150,
-            temperature=0.7,
-            top_p=0.95,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.95,
+        repetition_penalty=1.2
+    )
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response.split("### Response:")[-1].strip()
 
 # Streamlit UI
-st.set_page_config(page_title="Python Tutor Chatbot", page_icon="🐍")
 st.title("🐍 Python Tutor Chatbot")
-st.write("Ask me Python programming questions!")
+st.markdown("Ask me Python programming questions!")
 
-user_input = st.text_input("Your question:")
+user_input = st.text_area("Your question:")
 
-if user_input:
-    with st.spinner("Generating response..."):
-        response = chat(user_input)
+if st.button("Answer"):
+    if user_input.strip():
+        with st.spinner("Thinking..."):
+            answer = chat(user_input)
         st.markdown("**Answer:**")
-        st.markdown(response)
+        st.write(answer)
+    else:
+        st.warning("Please enter a question.")