Create app.py

app.py

@@ -0,0 +1,88 @@
+import os, transformers, peft, torch, gradio as gr
+
+base_model = "h2oai/h2ogpt-4096-llama2-13b"
+model = transformers.AutoModelForCausalLM.from_pretrained(
+    base_model,    
+    load_in_8bit=True,
+    torch_dtype=torch.float16
+)
+tokenizer = transformers.AutoTokenizer.from_pretrained(base_model)
+
+lora_model = "pcalhoun/Llama-2-13b-Conversations-With-Tyler-Swift"
+model = peft.PeftModel.from_pretrained(
+    model,
+    lora_model,
+    torch_dtype=torch.float16
+)
+
+def generate(prompt, extra_eos=[]):
+    inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
+    input_token_length = inputs.input_ids.shape[1]
+    outputs = model.generate(**inputs, max_length=4096)
+    text = tokenizer.batch_decode(outputs)[0]
+    return text
+
+def create_next_prompt(title_string,description_string="",conversation_messages=[]):
+    if not len(description_string):
+        conversation_messages = []
+    prompt = """<s>### CONVERSATIONS WITH TYLER SWIFT ###
+
+TITLE: """ + title_string.strip() + """
+
+DESCRIPTION:"""
+    if not len(description_string):
+        return prompt
+    else:
+        prompt += " "+description_string.replace("\n\n","\n").strip() + "\n\n"
+    if not len(conversation_messages):
+        prompt += "### TYLER SWIFT:"
+        return prompt
+    else:
+        for message_data in conversation_messages:
+            prompt += "### " + message_data['speaker'].upper() + ": " + message_data['message'].strip()
+            if message_data['speaker'].upper() == "TYLER SWIFT":
+                prompt += "</s><s>"
+            prompt += "\n"
+        if conversation_messages[-1]["speaker"].upper() != "TYLER SWIFT":
+            prompt += "### TYLER SWIFT:"
+        return prompt
+
+def deconstruct_returned_text(text):
+    #skip first line
+    text = "\n".join(text.split("\n")[1:]).strip()
+    title = text.split("\n")[0].replace("TITLE:","").strip()
+    text = "\n".join(text.split("\n")[1:]).strip()
+    description = text.split("\n\n")[0].replace("DESCRIPTION:","").strip()
+    text = "\n\n".join(text.split("\n\n")[1:]).strip()
+    conversation_text = text.replace("</s>", "").replace("<s>", "").split("<<")[0].strip()
+    return title,description,conversation_text
+
+def generate_next(title,description,conversation_text):
+    if not len(title):
+        title = "Set a Title First"
+        return title,description,conversation_text
+    conversation = []
+    for line in conversation_text.split("\n"):
+        if "CONVERSATIONS WITH TYLER SWIFT" in line:
+            continue
+        if line.startswith("###"):
+            speaker = line.split(":")[0].replace("###","").strip()
+            message = ":".join(line.split(":")[1:]).strip()
+            conversation.append({"speaker":speaker,"message":message.replace("</s>", "").replace("<s>", "").strip()})
+    prompt = create_next_prompt(title,description,conversation)
+    generated_text = generate(prompt)
+    print("GENERATED TEXT:",generated_text)
+    title,description,conversation_text = deconstruct_returned_text(generated_text)
+    return title,description,conversation_text
+
+
+with gr.Blocks() as demo:
+    title = gr.Textbox(label="Title")
+    description = gr.Textbox(label="Description")
+    conversation_text = gr.Textbox(label="Conversation")
+    generate_button = gr.Button()
+    generate_button.click(fn=generate_next, inputs=[title,description,conversation_text], outputs=[title,description,conversation_text])
+
+demo.launch()
+
+