Update app.py

app.py

@@ -1,15 +1,25 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load the text generation pipeline
-pipe = pipeline("text-generation")
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+model = AutoModelForCausalLM.from_pretrained("google/gemma-7b", device="cuda" if st.sidebar.checkbox("Use GPU", True) else "cpu")
 
-# Streamlit UI
-st.title("Demo with Hugging Face")
+# Function to generate text based on user input
+def generate_text(prompt):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
+    outputs = model.generate(input_ids, max_length=100)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
-prompt = st.text_area("Enter your text prompt:")
+# Streamlit app
+st.title("Text Generation with Google Gemma 7b")
 
+prompt = st.text_area("Enter your prompt here:", "")
 if st.button("Generate Text"):
-    # Generate text based on the prompt
-    generated_text = pipe(prompt, max_length=100, do_sample=True)[0]["generated_text"]
-    st.write(generated_text)
+    if prompt:
+        generated_text = generate_text(prompt)
+        st.write("Generated Text:")
+        st.write(generated_text)
+    else:
+        st.warning("Please enter a prompt.")