Update app.py

app.py

@@ -4,33 +4,31 @@ import streamlit as st
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM 
 
-'''
+
 tokenizer = AutoTokenizer.from_pretrained(
-  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b',
+  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b', cache_dir='./model_dir/',
   bos_token='[BOS]', eos_token='[EOS]', unk_token='[UNK]', pad_token='[PAD]', mask_token='[MASK]'
 )
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 model = AutoModelForCausalLM.from_pretrained(
-  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b',
+  'kakaobrain/kogpt', revision='KoGPT6B-ryan1.5b',cache_dir='./model_dir/',
   pad_token_id=tokenizer.eos_token_id,
   torch_dtype=torch.float16, low_cpu_mem_usage=False
 ).to(device=device, non_blocking=True)
 _ = model.eval()
-'''
+
 print("Model loading done!")
 
 def gpt(prompt):
-  return prompt
-  '''
   with torch.no_grad():
     tokens = tokenizer.encode(prompt, return_tensors='pt').to(device=device, non_blocking=True)
     gen_tokens = model.generate(tokens, do_sample=True, temperature=0.8, max_length=256)
     generated = tokenizer.batch_decode(gen_tokens)[0]
 
   return generated
-  '''
+  
     
 #prompts
 st.title("여러분들의 문장을 완성해줍니다. 🤖")