Update app.py

app.py

@@ -1,24 +1,32 @@
 import streamlit as st
-from transformers import pipeline, AutoProcessor, AutoModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import transformers
 import torch
-from diffusers import FluxPipeline
-from IPython.display import Audio
 from huggingface_hub import login
 import os
 
+# load huggingface
 api_key = os.getenv("ACCESS_TOKEN")
 login(token=api_key)
 
+# setup model
+model_id = "google/gemma-2-2b-it"
+dtype = torch.bfloat16
+chat = []
 
-# llama
-pipe = pipeline(task="text-generation", model="google/gemma-2-2b-it")
-messages = []
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=dtype,)
 
 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
-    messages.append({"role": "user", "content": prompt})
-    outputs = pipe(messages, max_length=80)
-    return outputs[0]["generated_text"][-1]["content"].strip()
+    chat.append({"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")
+    outputs = model.generate(input_ids=inputs.to(model.device), max_new_tokens=150)
+    return tokenizer.decode(outputs[0])
     
 st.title("Shakespeare Ai")
 st.write("A space made to allow people to create shakespeare like text!")