new stuff

app.py

@@ -1,7 +1,6 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
-from peft import PeftModel
 
 base_model_name = "chaseharmon/Rap-Mistral-Big"
 
@@ -14,7 +13,7 @@ def load_model():
         bnb_4bit_use_double_quant=False,
         bnb_4bit_compute_dtype="float16"
     )
-    
+
     model = AutoModelForCausalLM.from_pretrained(
         base_model_name,
         device_map='auto',
@@ -33,15 +32,24 @@ def load_tokenizer():
 
     return tokenizer
 
-
+def build_prompt(question):
+  prompt=f"[INST] {question} [/INST] "
+  return prompt
 
 model = load_model()
+model.eval()
+
 tokenizer = load_tokenizer
 
 st.title("Rap Verse Generation V1 Demo")
 st.header("Supported Artists")
 st.write("Lupe Fiasco, Common, Jay-Z, Yasiin Bey, Ab-Soul, Rakim")
 
-prompt = st.chat_input("Write a verse in the style of Lupe Fiasco")
+question = st.chat_input("Write a verse in the style of Lupe Fiasco")
+if question:
+    prompt = build_prompt(question)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    model_inputs = inputs.to('cuda')
+    generated_ids = model.generate(**model_inputs, max_new_tokens=300, do_sample=True, pad_token_id=tokenizer.eos_token_id)
+    decoded_output = tokenizer.batch_decode(generated_ids)
 
-st.write(prompt)