Update app.py

app.py

@@ -1,9 +1,12 @@
 import streamlit as st
 
 def summarize_function(notes):
-    gen_text = pipe(notes, max_length=(len(notes.split(' '))*2*1.225), temperature=0.8, num_return_sequences=1, top_p=0.2)[0]['generated_text']
+    gen_text = pipe(notes, max_length=(len(notes.split(' '))*2*1.225), temperature=0.8, num_return_sequences=1, top_p=0.2)[0]['generated_text'][len(notes):]
+    for i in range(len(gen_text)):
+        if gen_text[-i-8:].startswith('[Notes]:'):
+            gen_text = gen_text[:-i-8]
     st.write('Summary: ')
-    return gen_text[len(notes):]
+    return gen_text
 
 st.markdown("<h1 style='text-align: center; color: #489DDB;'>GPT Clinical Notes Summarizer 0.1v</h1>", unsafe_allow_html=True)
 st.markdown("<h6 style='text-align: center; color: #489DDB;'>by Bryan Mildort</h1>", unsafe_allow_html=True)
@@ -17,8 +20,8 @@ import torch
 # device_map = infer_auto_device_map(model, dtype="float16")
 # st.write(device_map)
 
-model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes_summary", low_cpu_mem_usage=True).cuda()
-tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes_summary")
+model = AutoModelForCausalLM.from_pretrained("bryanmildort/gpt_neo_notes", low_cpu_mem_usage=True).cuda()
+tokenizer = AutoTokenizer.from_pretrained("bryanmildort/gpt_neo_notes")
 # model = model.to(device)
 
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)