Update app.py

app.py

@@ -1,38 +1,23 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextStreamer
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 
-MODEL_PATH = "./tinyllama-python-tutor-lora"
-
-st.title("TinyLLaMA Python Tutor 💬")
+st.title("TinyLLaMA Python Tutor (LoRA)")
 
 @st.cache_resource
 def load_model():
-    bnb_config = BitsAndBytesConfig(
-        load_in_8bit=True,
-        llm_int8_threshold=6.0,
-        llm_int8_skip_modules=None,
-        llm_int8_enable_fp32_cpu_offload=True
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, use_fast=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_PATH,
-        quantization_config=bnb_config,
-        device_map="auto"
-    )
+    base_model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+    tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
+    model = PeftModel.from_pretrained(base_model, "lora_adapter")
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-prompt = st.text_area("Ask me about Python:", height=200)
+user_input = st.text_area("Ask me a Python coding question:")
 
-if st.button("Generate Response"):
-    if prompt.strip():
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-        output = model.generate(**inputs, max_new_tokens=200, do_sample=True)
-        response = tokenizer.decode(output[0], skip_special_tokens=True)
-        st.write("### Response")
-        st.success(response)
-    else:
-        st.warning("Please enter a prompt!")
+if st.button("Generate Answer"):
+    inputs = tokenizer(user_input, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=150)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    st.write("**Answer:**", response)