Update app.py

app.py

@@ -4,6 +4,11 @@ import nltk
 from nltk import sent_tokenize
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import json
+
+from sentence_transformers import SentenceTransformer
+
+sentence_transformer_model = SentenceTransformer("sentence-transformers/all-roberta-large-v1")
+
 nltk.download('punkt')
 with open('testbook.json') as f:
     test_book = json.load(f)
@@ -63,7 +68,7 @@ def chunking(book_text):
                     next_token_len += token_lens[i+t] 
                     next_pseudo_segment += sentences[i+t] 
                 
-            embs = model.encode([current_segment, next_pseudo_segment, sentences[i]]) # current, next, sent
+            embs = sentence_transformer_model.encode([current_segment, next_pseudo_segment, sentences[i]]) # current, next, sent
             if cos_similarity(embs[1],embs[2]) > cos_similarity(embs[0],embs[2]):
                 segments.append(current_segment)
                 current_segment = sentences[i]