Update app.py

app.py

@@ -1,56 +1,48 @@
 import torch
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import pipeline
 import streamlit as st
 
+# Set up offload directory for CPU offloading
+offload_dir = "./offload"
+
 # Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("TinyLLaMA/TinyLLaMA-1.1B-Chat-v1.0")
+tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+
+# Load base model with CPU offloading config
+bnb_config = BitsAndBytesConfig(load_in_4bit=True)
 
-# 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"
+    "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    quantization_config=bnb_config,
+    device_map="auto",  # required for offloading
+    offload_folder=offload_dir  # this is the key line
 )
 
-# Load LoRA adapter
-model = PeftModel.from_pretrained(base_model, "lora_adapter")
+# Load your LoRA adapter
+model = PeftModel.from_pretrained(base_model, "lora_adapter", device_map="auto")
 
-# Set title
-st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
-st.markdown("Ask me any **Python programming** question:")
+# Text generation pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# User input
-user_question = st.text_input("Your question")
+# Streamlit UI
+st.title("🧠 TinyLLaMA Python Tutor (LoRA)")
+st.write("Ask me any Python programming question:")
 
-if user_question:
-    with st.spinner("Thinking..."):
+user_input = st.text_input("Your question")
 
-        # 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."
+if user_input:
+    # Check if question is Python-related
+    if "python" in user_input.lower() or "list" in user_input.lower() or "def " in user_input.lower() or "tuple" in user_input.lower() or "function" in user_input.lower():
+        prompt = f"""You are a helpful Python tutor. Answer only Python programming questions. 
+Respond clearly with examples. Avoid repeating the question.
 
-Question: {user_question}
+Question: {user_input}
 Answer:"""
-
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-        output = model.generate(
-        **inputs,
-        max_new_tokens=512,   # allow longer answers
-        do_sample=True,
-        top_p=0.9,
-        temperature=0.7,
-        repetition_penalty=1.1
-
-        )
-
-        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()
-
+        
+        response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
+        answer = response.split("Answer:")[-1].strip()
         st.markdown(f"💬 **Answer:**\n\n{answer}")
+    else:
+        st.markdown("❌ Sorry, I can only answer Python programming questions.")