Update app.py

app.py

@@ -1,31 +1,21 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import transformers
+from transformers import pipeline, AutoProcessor, AutoModel
 import torch
 from huggingface_hub import login
 import os
 
-# load huggingface
+#login
 api_key = os.getenv("ACCESS_TOKEN")
 login(token=api_key)
 
-# setup model
-model_id = "google/gemma-2-2b-it"
-dtype = torch.bfloat16
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
-model.to(device)
-model.eval()
 
-def poet(text):
-    prompt = 'Make 25 lines, it has to be absolutely 25 lines of text no less no exception, of shakespeare based on this prompt: ' + text
-    chat = [{"role": "user", "content": prompt}]
-    prompt = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
-    inputs = inputs.to(device)
-    outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=350)
-    return tokenizer.decode(outputs[0])
+# gemma
+pipe = pipeline(task="text-generation", model="google/gemma-2-2b-it")
+
+def poet(prompt):
+    messages = [{"role": "user", "content": prompt}]
+    outputs = pipe(messages, max_length=1000)
+    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!")