# 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-3-1b-it"
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.")