Update app.py

app.py

@@ -1,16 +1,20 @@
-"""app.py"""
-
 import streamlit as st
-from transformers import pipeline, GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_name = "Llama-2-7b-finetuned-with-QLoRa"
+
+# Load model and tokenizer
+@st.cache_resource
+def load_model_and_tokenizer(model_name):
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
 
-# Load pre-trained GPT-2 model and tokenizer
-model_name = "gpt2"
-model = GPT2LMHeadModel.from_pretrained(model_name)
-tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+model, tokenizer = load_model_and_tokenizer(model_name)
 
-# Define function to generate blog post
-def generate_blogpost(topic):
-    input_text = f"Blog post about {topic}:"
+# Function to generate response
+def generate_response(topic):
+    input_text = f"Response about {topic}:"
     input_ids = tokenizer.encode(input_text, return_tensors="pt")
 
     # Generate text
@@ -22,16 +26,17 @@ def generate_blogpost(topic):
 
 # Streamlit app
 def main():
-    st.title("Blog Post Generator")
+    st.title("Llama 2 Fine-Tuned Demo with QLoRa")
 
     # Sidebar input for topic
-    topic = st.sidebar.text_input("Enter topic for the blog post", "a crazy person driving a car")
+    topic = st.sidebar.text_input("Enter your topic", "a crazy person driving a car")
 
     # Generate button
-    if st.sidebar.button("Generate Blog Post"):
-        blogpost = generate_blogpost(topic)
-        st.subheader(f"Generated Blog Post on {topic}:")
-        st.write(blogpost)
+    if st.sidebar.button("Generate Response"):
+        with st.spinner("Generating response..."):
+            response = generate_response(topic)
+        st.subheader(f"Generated response on '{topic}':")
+        st.write(response)
 
 if __name__ == "__main__":
     main()