Update app.py

app.py

@@ -6,69 +6,31 @@ import time
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import streamlit as st
+
+from transformers import AutoModelForSeq2SeqLM
+import torch
+
 from threading import Thread
 os.environ["COQUI_TOS_AGREED"] = "1"
 os.environ["TRAINER_TELEMETRY"]= "0"
 # Constants
-MODEL_LIST = ["mistralai/Mistral-Nemo-Instruct-2407"]
+
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL = os.environ.get("MODEL_ID", "mistralai/Mistral-Nemo-Instruct-2407")
+
 # Set the device to GPU or CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-    ignore_mismatched_sizes=True
-).to(device)
-
-# CSS styles
-CSS = """
-<style>
-h1 {
-    text-align: center;
-}
-</style>
-"""
-
-
-# Function to handle chat
-def stream_chat(message, history, temperature=0.3, max_new_tokens=1024, top_p=1.0, top_k=20, penalty=1.2):
-    conversation = []
-    for prompt, answer in history:
-        conversation.extend([
-            {"role": "user", "content": prompt}, 
-            {"role": "assistant", "content": answer},
-        ])
-    
-    conversation.append({"role": "user", "content": message})
-    
-    input_text = tokenizer.apply_chat_template(conversation, tokenize=False)
-    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
-    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
-    
-    generate_kwargs = dict(
-        input_ids=inputs, 
-        max_new_tokens=max_new_tokens,
-        do_sample=temperature != 0,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        streamer=streamer,
-        pad_token_id=10,
-    )
-
-    with torch.no_grad():
-        thread = Thread(target=model.generate, kwargs=generate_kwargs)
-        thread.start()
-        
-    buffer = ""
-    for new_text in streamer:
-        buffer += new_text
-        yield buffer
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("merve/chatgpt-prompts-bart-long")
+model = AutoModelForSeq2SeqLM.from_pretrained("merve/chatgpt-prompts-bart-long", from_tf=True).to("cuda" if torch.cuda.is_available() else "cpu")
+
+# Function to generate the prompt based on the persona
+def generate(prompt):
+    batch = tokenizer(prompt, return_tensors="pt").to(model.device)
+    generated_ids = model.generate(batch["input_ids"], max_new_tokens=150)
+    output = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+    return output[0]
+
 
 #st.set_page_config(layout="wide")
 # Load custom CSS to integrate Bootstrap, Font Awesome, and Google Fonts
@@ -150,34 +112,17 @@ with left:
     st.markdown('''<h3><i class="fa fa-pencil"></i> Form 1</h3>''', unsafe_allow_html=True)
 
     # Box 2: Form 1
-    # Chat history
-    if 'history' not in st.session_state:
-        st.session_state['history'] = []
-
-    # Chat input
-    user_input = st.text_input("Your Message:", "")
-
-    # Chat parameters
-    with st.expander("⚙️ Parameters"):
-        temperature = st.slider("Temperature", 0.0, 1.0, 0.3)
-        max_new_tokens = st.slider("Max new tokens", 128, 8192, 1024)
-        top_p = st.slider("Top p", 0.0, 1.0, 1.0)
-        top_k = st.slider("Top k", 1, 20, 20)
-        penalty = st.slider("Repetition penalty", 0.0, 2.0, 1.2)
-
-    # Handle the chat logic
-    if st.button("Send"):
-        if user_input:
-            response = stream_chat(user_input, st.session_state['history'], temperature, max_new_tokens, top_p, top_k, penalty)
-            st.session_state['history'].append((user_input, next(response)))
-            for new_text in response:
-                st.session_state['history'][-1] = (user_input, new_text)
-                st.experimental_rerun()
-
-    # Display chat history
-    for prompt, answer in st.session_state['history']:
-        st.write(f"**User:** {prompt}")
-        st.write(f"**Mistral-Nemo:** {answer}")
+    persona = st.text_input("Input a persona, e.g. photographer", value="photographer")
+
+    # Button to trigger generation
+    if st.button("Generate Prompt"):
+        if persona:
+            with st.spinner("Generating..."):
+                result = generate(persona)
+            st.text_area("Generated Prompt", value=result, height=200)
+        else:
+            st.error("Please enter a persona to generate a prompt.")
+
     
 
 with right: