Update app.py

app.py

@@ -4,9 +4,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 @st.cache(allow_output_mutation=True)
-def opt_model(prompt, num_sequences = 1, max_length = 50):  
-    model = AutoModelForCausalLM.from_pretrained("facebook/opt-30b", torch_dtype=torch.float16).cuda()
-    tokenizer = AutoTokenizer.from_pretrained("facebook/opt-30b", use_fast=False)
+model = AutoModelForCausalLM.from_pretrained("facebook/opt-30b", torch_dtype=torch.float16).cuda()
+tokenizer = AutoTokenizer.from_pretrained("facebook/opt-30b", use_fast=False)
+
+@st.cache(allow_output_mutation=True)
+def opt_model(prompt, model, tokenizer, num_sequences = 1, max_length = 50):  
     input_ids = tokenizer(prompt, return_tensors="pt").input_ids.cuda()
     generated_ids = model.generate(input_ids, num_return_sequences=num_sequences, max_length=max_length)
     answer = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
@@ -15,7 +17,7 @@ def opt_model(prompt, num_sequences = 1, max_length = 50):
 prompt= st.text_area('Your prompt here',
  '''Hello, I'm am conscious and''')
 
-answer = opt_model(prompt)
+answer = opt_model(prompt, model, tokenizer,)
 #lst = ['ciao come stai sjfsbd dfhsdf  fuahfuf  feuhfu wefwu ']
 lst = ' '.join(answer)