Update app.py

app.py

@@ -1,29 +1,24 @@
 import torch
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 from transformers import pipeline
 import streamlit as st
-
-# Set up offload directory for CPU offloading
-offload_dir = "./offload"
+import os
 
 # Load tokenizer
 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",
-    quantization_config=bnb_config,
-    device_map="auto",  # required for offloading
-    offload_folder=offload_dir  # this is the key line
+    device_map="auto",
+    torch_dtype=torch.float32
 )
 
-# Load your LoRA adapter
+# Load LoRA adapter
 model = PeftModel.from_pretrained(base_model, "lora_adapter", device_map="auto")
 
-# Text generation pipeline
+# Load pipeline
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Streamlit UI
@@ -33,16 +28,13 @@ st.write("Ask me any Python programming question:")
 user_input = st.text_input("Your question")
 
 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.
+    if "python" in user_input.lower() or "list" in user_input.lower() or "tuple" in user_input.lower() or "def " in user_input.lower() or "class" in user_input.lower():
+        prompt = f"""You are a helpful and friendly Python tutor. Only answer Python programming questions. Be clear and concise.
 
 Question: {user_input}
 Answer:"""
-        
-        response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
+        response = pipe(prompt, max_new_tokens=256, temperature=0.7, do_sample=True)[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.")
+        st.warning("❌ Sorry, I can only answer Python programming questions.")