Update app.py

app.py

@@ -27,10 +27,11 @@ def get_model():
     #model = AutoModelForCausalLM.from_pretrained("facebook/opt-125m")
     tokenizer = AutoTokenizer.from_pretrained("BigSalmon/InformalToFormalLincoln63Paraphrase")
     model = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln68Paraphrase")
+    model2 = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln63Paraphrase")
     #model = AutoModelForCausalLM.from_pretrained("BigSalmon/InformalToFormalLincoln63Paraphrase")
-    return model, tokenizer
+    return model, model2, tokenizer
     
-model, tokenizer = get_model()
+model, model2, tokenizer = get_model()
 
 st.text('''For Prompt Templates: https://huggingface.co/BigSalmon/InformalToFormalLincoln63Paraphrase''')
 
@@ -65,10 +66,41 @@ def run_generate(text, bad_words):
     e = e.replace(text, "")
     yo.append(e)
   return yo
+  
+def run_generate2(text, bad_words):
+  yo = []
+  input_ids = tokenizer.encode(text, return_tensors='pt')
+  res = len(tokenizer.encode(text))
+  bad_words = bad_words.split()
+  bad_word_ids = []
+  for bad_word in bad_words: 
+    bad_word = " " + bad_word
+    ids = tokenizer(bad_word).input_ids
+    bad_word_ids.append(ids)
+  sample_outputs = model2.generate(
+    input_ids,
+    do_sample=True, 
+    max_length= res + lengths, 
+    min_length = res + lengths, 
+    top_k=50,
+    temperature=temp,
+    num_return_sequences=number_of_outputs,
+    bad_words_ids=bad_word_ids
+  )
+  for i in range(number_of_outputs):
+    e = tokenizer.decode(sample_outputs[i])
+    e = e.replace(text, "")
+    yo.append(e)
+  return yo
+ 
 with st.form(key='my_form'):
     text = st.text_area(label='Enter sentence', value=first)
     submit_button = st.form_submit_button(label='Submit')
     submit_button2 = st.form_submit_button(label='Submit Log Probs')
+    
+    submit_button3 = st.form_submit_button(label='Submit Other Model')
+    submit_button4 = st.form_submit_button(label='Submit Log Probs Other Model')
+    
     if submit_button:
       translated_text = run_generate(text, bad_words)
       st.write(translated_text if translated_text else "No translation found")
@@ -84,4 +116,19 @@ with st.form(key='my_form'):
         probabilities = torch.nn.functional.softmax(logits)
         best_logits, best_indices = logits.topk(logs_outputs)
         best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]      
+        st.write(best_words)
+    if submit_button3:
+      translated_text = run_generate2(text, bad_words)
+      st.write(translated_text if translated_text else "No translation found")
+    if submit_button4:
+        text2 = str(text)
+        print(text2)
+        text3 = tokenizer.encode(text2)
+        myinput, past_key_values = torch.tensor([text3]), None
+        myinput = myinput
+        logits, past_key_values = model2(myinput, past_key_values = past_key_values, return_dict=False)
+        logits = logits[0,-1]
+        probabilities = torch.nn.functional.softmax(logits)
+        best_logits, best_indices = logits.topk(logs_outputs)
+        best_words = [tokenizer.decode([idx.item()]) for idx in best_indices]      
         st.write(best_words)