Update app.py

app.py

@@ -1,20 +1,33 @@
+import streamlit as st
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-
 import os
 from huggingface_hub import login
 
+# Login to Hugging Face using the token
 read_key = os.environ.get("Read_Token")
 login(token=read_key)
 
+# Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
 model = AutoModelForCausalLM.from_pretrained(
     "google/gemma-2b-it",
     torch_dtype=torch.bfloat16
 )
 
-input_text = "Write me a poem about Machine Learning."
-input_ids = tokenizer(input_text, return_tensors="pt")
+# Streamlit UI
+st.title("Gemma-2B-IT Text Generator")
+st.write("Enter a prompt to generate text using the model.")
+
+# Input prompt from the user
+input_text = st.text_input("Prompt", "Write me a poem about Machine Learning.")
 
-outputs = model.generate(**input_ids)
-print(tokenizer.decode(outputs[0]))
+# Generate text when the button is pressed
+if st.button("Generate"):
+    input_ids = tokenizer(input_text, return_tensors="pt")
+    outputs = model.generate(**input_ids)
+    generated_text = tokenizer.decode(outputs[0])
+    
+    # Display the generated text
+    st.write("Generated Text:")
+    st.text_area("Output", generated_text, height=300)