Update app.py

app.py

@@ -1,23 +1,24 @@
 import gradio as gr
 from transformers import pipeline
 
-
-pipe = pipeline("text-generation", model="distilgpt2")
+# Initialize the text generation pipeline with the Gemma 2-2B IT model
+pipe = pipeline("text-generation", model="google/gemma-2-2b-it")
 
 def generate_response(user_input):
-    formatted_input = f"Human: {user_input}\nAI:"
-    response = pipe(formatted_input, max_length=150, num_return_sequences=1)
+    # Generate text based on the user's input
+    response = pipe(user_input, max_length=100, num_return_sequences=1)
+    # Extract the generated text
     generated_text = response[0]['generated_text']
-    ai_response = generated_text.split('AI:')[-1].strip()
-    return ai_response
-
+    return generated_text
 
+# Create the Gradio interface
 interface = gr.Interface(
     fn=generate_response,
-    inputs=gr.Textbox(label="shitty mesage:", lines=2, placeholder="Type your message here..."),
+    inputs=gr.Textbox(label="prompt:", lines=2, placeholder="prompt"),
     outputs="text",
-    title="distilgpt2",
-    description="distel gppt 20"
+    title="Gemma",
+    description="Prompt gemma-2b"
 )
 
+# Launch the Gradio app
 interface.launch()