updated from benny

app.py

@@ -1,8 +1,14 @@
+
+​
+Lillian McAboy
+​
+
 import streamlit as st
-import torch, os 
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
 
-model = 'google/gemma-1.1-7b-it'
+model_id = 'google/gemma-2b'
 
 formal_options = [
     'More aggressive',
@@ -19,44 +25,29 @@ formal_options = [
     'Like Eric Cartman'
 ]
 
-formality_name = st.selectbox("Transform my text to be: ", formal_options)
-
-@st.cache(allow_output_mutation=True)
-def get_tokenizer(model_name):
-    from transformers import AutoTokenizer
-    return AutoTokenizer.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, device_map='auto', torch_dtype=torch.bfloat16)
 
-@st.cache(allow_output_mutation=True)
-def get_model(model_name):
-    import torch
-    from transformers import AutoModelForCausalLM
-    mymodel = AutoModelForCausalLM.from_pretrained(model, device_map='auto', torch_dtype=torch.bfloat16)
-    print(f"Loaded model, {mymodel.num_parameters():,d} parameters.")
-    return mymodel
-original_text = st.text_area("Original Document", "This is my text that I want to change.")
-new_text = st.text_area("New Document", help="Leave this blank, this will be updated with your new document.")
+formality_name = st.selectbox("Transform my text to be: ", formal_options)
+original_text = st.text_area("Input", "This is my text that I want to change.")
+new_text = st.text_area("Output", key="output", help="Leave this blank, this will be updated with your new document.")
 
 def get_spans_local(formality_name, original_text, new_text):
-    #import torch
-    
-    tokenizer = get_tokenizer(model)
-    mymodel = get_model(model)
-    
     messages = [
         {
             "role": "user",
             "content": f"{formality_name}\n\n{original_text}",
         },
     ]
-    
+   
     tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")[0]
-    
+   
     new_text_ids = tokenizer(new_text, return_tensors='pt')['input_ids'][0]
     joined_ids = torch.cat([tokenized_chat, new_text_ids[1:]])
-    
+   
     with torch.no_grad():
         logits = model(joined_ids[None].to(model.device)).logits[0].cpu()
-        
+       
     spans = []
     length_so_far = 0
     for idx in range(len(tokenized_chat), len(joined_ids)):
@@ -77,25 +68,23 @@ def get_spans_local(formality_name, original_text, new_text):
     return spans
 
 def summarize_text(text):
-    input_ids = tokenizer.encode(text, return_tensors="pt", max_length=1024, truncation=True)
+    input_ids = tokenizer.encode(text, return_tensors="pt", max_length=1024, truncation=True).to("cuda")
     summary_ids = model.generate(input_ids, max_length=150, num_beams=2, early_stopping=True)
     summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return summary
 
 
-def get_highlights_api(prompt, doc, updated_doc):
-    # Make a request to the API. prompt and doc are query parameters:
-    # https://tools.kenarnold.org/api/highlights?prompt=Rewrite%20this%20document&doc=This%20is%20a%20document
-    # The response is a JSON array
-    import requests
-    response = requests.get("https://tools.kenarnold.org/api/highlights", params=dict(prompt=prompt, doc=doc, updated_doc=updated_doc))
-    return response.json()['highlights']
+# def get_highlights_api(prompt, doc, updated_doc):
+#     # Make a request to the API. prompt and doc are query parameters:
+#     # https://tools.kenarnold.org/api/highlights?prompt=Rewrite%20this%20document&doc=This%20is%20a%20document
+#     # The response is a JSON array
+#     import requests
+#     response = requests.get("https://tools.kenarnold.org/api/highlights", params=dict(prompt=prompt, doc=doc, updated_doc=updated_doc))
+#     return response.json()['highlights']
+
+
+# # spans = get_spans_local(formality_name, original_text, new_text)
+# # st.write(spans)
 
-if __name__ == "__main__":
-    spans = get_spans_local(formality_name, original_text, new_text)
-    st.write(spans)
-    
-    if st.button("Summarize My Document"):
-        summary = summarize_text(original_text)
-        summary_text = summary
-        st.write(summary_text)
+if st.button("Summarize My Document"):
+    new_text = st.text_area("Output", value=summarize_text(original_text), key="output", help="Leave this blank, this will be updated with your new document.")