Update app.py

app.py

@@ -11,16 +11,15 @@ login(token=api_key)
 
 
 # llama
-pipe = pipeline(task="text-generation", model="google/gemma-2-2b-jpn-it")
-messages = [
-    {"role": "system", "content": "You are a text generator that writes Shakespeare given a prompt, the text you write needs be 25 lines long. and cannot be none, you must write a poem that is 25 lines long."},
-]
+pipe = pipeline(task="text-generation", model="google/gemma-2-2b-it")
+messages = []
 
 def poet(text):
-    messages.append({"role": "user", "content": text})
+    prompt = 'Make 25 lines of shakespeare based on: ' + text
+    messages.append({"role": "user", "content": prompt})
     outputs = pipe(messages, max_length=80)
-    return outputs[0]["generated_text"].strip()
-
+    return outputs[0]["generated_text"][-1]["content"].strip()
+    
 st.title("Shakespeare Ai")
 st.write("A space made to allow people to create shakespeare like text with images!")