basic streamlit interface

.gitignore

@@ -0,0 +1,2 @@
+.venv/
+__pycache__/

app.py

@@ -0,0 +1,38 @@
+
+ascii_art = r"""
+      _               _            _    
+  ___| |_ __ _ _ __  | |_ _ __ ___| | __
+ / __| __/ _` | '__| | __| '__/ _ \ |/ /
+ \__ \ || (_| | |    | |_| | |  __/   < 
+ |___/\__\__,_|_|     \__|_|  \___|_|\_\
+
+"""
+
+from model import LLM
+import streamlit as st
+import streamlit_scrollable_textbox as stx
+
+st.text(ascii_art)
+
+with st.spinner("Please wait... loading model"):
+    llm = LLM()
+
+demo_text = """DATA: The ship has gone into warp, sir.
+RIKER: Who gave the command?
+DATA: Apparently no one. Helm and navigation controls are not functioning. Our speed is now warp seven-point-three and holding.
+PICARD: Picard to Engineering. Mister La Forge, what's going on down there?
+"""
+text = st.text_area("First few lines of the script:", demo_text)
+
+col1, col2 = st.columns(2)
+
+with col1:
+    temp = st.slider('Temperature', 0, 1, 1)
+    max_len = st.number_input('Max length', min_value=1, max_value=2048, value=512)
+
+with col2:
+    top_p = st.slider('p', 0, 1, 0.95)
+    top_k = st.slider('k', 1, 100, 50)
+
+with st.spinner("Generating text..."):
+    stx.scrollableTextbox(llm.generate(text, max_len, temp, top_k, top_p), height=400)

model.py

@@ -0,0 +1,23 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class LLM:
+    def __init__(self):
+        self.model = AutoModelForCausalLM.from_pretrained('progs2002/star-trek-tng-script-generator')
+        self.tokenizer = AutoTokenizer.from_pretrained('progs2002/star-trek-tng-script-generator')
+    
+    def generate(self, text, max_len=512, temp=1, k=50, p=0.95):
+        encoded_prompt = self.tokenizer.encode(text, add_special_tokens=False, return_tensors="pt")
+
+        output_tokens = self.model.generate(
+            input_ids = encoded_prompt,
+            max_length = max_len,
+            do_sample=True,
+            num_return_sequences=1,
+            pad_token_id=self.model.config.eos_token_id,
+            temperature=temp,
+            top_k=k,
+            top_p=p
+        )
+
+        text_out = self.tokenizer.decode(output_tokens[0], clean_up_tokenization_spaces=True, skip_special_tokens=True)
+        return text_out

requirements.txt

@@ -0,0 +1,3 @@
+transformers
+torch
+streamlit-scrollable-textbox