Update app.py

app.py

@@ -1,29 +1,41 @@
-# prompt: create a sreamlit app on finetuned llm
-
-import streamlit as st
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-
-# Load the fine-tuned model and tokenizer
-model_path = "fine_tuned_gemma_3_1b"  # Replace with the actual path to your model directory
-tokenizer = AutoTokenizer.from_pretrained(model_path)
-model = AutoModelForCausalLM.from_pretrained(model_path)
-
-# Create a text generation pipeline
-text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-st.title("Fine-tuned Gemma 3.1B LLM")
-
-# Create a text input box for the user
-user_input = st.text_area("Enter your prompt:")
-
-if st.button("Generate Text"):
-    if user_input:
-        # Generate text based on user input
-        output = text_generator(user_input, max_length=150, num_return_sequences=1)
-        generated_text = output[0]['generated_text']
-
-        # Display the generated text
-        st.write("Generated Text:")
-        st.write(generated_text)
-    else:
+# prompt: create a sreamlit app on finetuned llm
+
+import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
+
+# Load Base Gemma Model (Required for Adapter)
+base_model_path = "google/gemma-1b"
+model = AutoModelForCausalLM.from_pretrained(
+    base_model_path,
+    load_in_4bit=True,          # Efficient memory usage
+    device_map="auto"           # Automatically maps to GPU if available
+)
+
+# Load LoRA Adapter
+model = PeftModel.from_pretrained(model, "fine_tuned_gemma_3_1b/")
+
+# Load the fine-tuned model and tokenizer
+# model_path = "fine_tuned_gemma_3_1b"  # Replace with the actual path to your model directory
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+# model = AutoModelForCausalLM.from_pretrained(model_path)
+
+# Create a text generation pipeline
+text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+st.title("Fine-tuned Gemma 3.1B LLM")
+
+# Create a text input box for the user
+user_input = st.text_area("Enter your prompt:")
+
+if st.button("Generate Text"):
+    if user_input:
+        # Generate text based on user input
+        output = text_generator(user_input, max_length=150, num_return_sequences=1)
+        generated_text = output[0]['generated_text']
+
+        # Display the generated text
+        st.write("Generated Text:")
+        st.write(generated_text)
+    else:
         st.warning("Please enter a prompt.")