Borgesian caching added, truncation fixed

pages/Borgesian.py

@@ -3,16 +3,21 @@ import transformers
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
-borgesian = GPT2LMHeadModel.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2', output_attentions = False, output_hidden_states = False)
-borgesian.load_state_dict(torch.load('borgesian_weights.pt', map_location=torch.device('cpu')))
-tokenizer = GPT2Tokenizer.from_pretrained("sberbank-ai/rugpt3small_based_on_gpt2")
+@st.cache_resource
+def load_model():
+    borgesian = GPT2LMHeadModel.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2', output_attentions = False, output_hidden_states = False)
+    borgesian.load_state_dict(torch.load('borgesian_weights.pt', map_location=torch.device('cpu')))
+    tokenizer = GPT2Tokenizer.from_pretrained("sberbank-ai/rugpt3small_based_on_gpt2")
+    return borgesian, tokenizer
+
+borgesian, tokenizer = load_model()
 borgesian.to('cpu')
 borgesian.eval()
 
 def generate_response(text, temperature, length, top_p):
     input_ids = tokenizer.encode(text, return_tensors="pt")
     with torch.no_grad():
-        out = borgesian.generate(input_ids, do_sample=True, num_beams=2, temperature=float(temperature), top_p=float(top_p), max_length=length)
+        out = borgesian.generate(input_ids, do_sample=True, num_beams=2, temperature=float(temperature), top_p=float(top_p), max_length=length, truncate=".")
     generated_text = list(map(tokenizer.decode, out))[0]
     st.write(generated_text)