Update app.py

app.py

@@ -1,46 +1,59 @@
-import transformers
-import streamlit as st
-from transformers import AutoTokenizer, AutoModelWithLMHead
-
-tokenizer = AutoTokenizer.from_pretrained("gpt2-large")
-@st.cache_data
-def load_model(model_name):
-model = AutoModelWithLMHead.from_pretrained(model_name)
-return model
-
-model = load_model("gpt2-large")
-
-def infer(sent, max_length, temperature, top_k, top_p):
-input_ids = tokenizer.encode(sent, return_tensors="pt")
-output_sequences = model.generate(
-input_ids=input_ids,
-max_length=max_length,
-temperature=temperature,
-top_k=top_k,
-top_p=top_p,
-do_sample=True,
-num_return_sequences=1
-)
-
-return output_sequences
-
-default_value = "You: Ask me anything!"
+import os
+import sys
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from PIL import Image
+
+sys.path.append(".")
+os.environ["HUGGINGFACE_HTTPS_PROXY"] = "" # Disable HTTPS proxy when not required
+
+tokenizer = AutoTokenizer.from_pretrained("./model_dir")
+model = AutoModelForSeq2SeqLM.from_pretrained("./model_dir", device_map="auto").half()
+device = next(iter(model.parameters())).device
+
+def generate_response(input_text):
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(
+        input_ids,
+        max_length=512,
+        num_beams=5,
+        early_stopping=True,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+        length_penalty=1.0,
+        no_repeat_ngram_size=2,
+        min_length=10,
+        temperature=0.9,
+    )
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
-#prompts
-st.title("Chat with GPT-2 💬")
-st.write("GPT-2 is a large transformer-based language model with 1.5 billion parameters. It is trained to predict the next word in a sentence, given all of the previous words. This makes it great for text generation and for answering questions about the text it's given.")
-
-messages = [{"role": "system", "content": "You are a helpful assistant."}]
-
-user_input = st.text_input("You:", default_value)
-if user_input:
-messages.append({"role": "user", "content": user_input})
-
-output_sequences = infer(user_input, max_length=100, temperature=0.7, top_k=40, top_p=0.9)
-generated_sequence = output_sequences[0].tolist()
-generated_text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
-
-messages.append({"role": "assistant", "content": generated_text})
+import streamlit as st
 
-for message in messages:
-st.write(f"{message['role']}: {message['content']}")
+st.set_page_config(layout="wide")
+col1, col2 = st.beta_columns((3, 1))
+
+with open('style.css') as f:
+    st.markdown(f'<style>{f.read()}<\style>', unsafe_allow_html=True)
+
+with col1:
+    user_prompt = st.text_area("You:", "", height=50)
+    
+if 'generated' not in st.session_state:
+    st.session_state['generated'] = ''
+
+if len(user_prompt) > 0 and st.button('Send'):
+    response = generate_response(user_prompt)
+    st.write('<span style="font-weight:bold;">Assistant:</span>
+' + response, unsafe_allow_html=True)
+    st.session_state['generated'] += '\n\n<span style="font-weight:bold;">User:</span>'+'\n'+ user_prompt + '\n'
+    st.session_state['generated'] += '<span style="font-weight:bold;">Assistant:</span>\n' + response
+
+if 'generated' in st.session_state:
+    message = st.session_state['generated'].replace('\n', '
+').replace('<span style="font-weight: bold;">User:</span>', '&uarr;').replace('<span style="font-weight: bold;">Assistant:</span>', '')
+    st.markdown(message, unsafe_allow_html=True)
+
+# Upload logo
+logo = Image.open("your_logo.png")
+st.sidebar.image(logo, width=160)