Danil commited on
Commit
0bea55f
Β·
1 Parent(s): 7c0ff8b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -1,5 +1,5 @@
1
  import streamlit as st
2
- from transformers import GPT2LMHeadModel, GPT2Tokenizer
3
 
4
  @st.cache(allow_output_mutation=True)
5
  def load_model():
@@ -7,9 +7,9 @@ def load_model():
7
  Loads the model and tokenizer from the local directory.
8
  :return: A list containing the model and the tokenizer.
9
  '''
10
- model_name = 'WIP'
11
- tokenizer = GPT2Tokenizer.from_pretrained(model_name)
12
- model = GPT2LMHeadModel.from_pretrained(model_name)
13
  return [model, tokenizer]
14
 
15
  st.set_page_config(
 
1
  import streamlit as st
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
 
4
  @st.cache(allow_output_mutation=True)
5
  def load_model():
 
7
  Loads the model and tokenizer from the local directory.
8
  :return: A list containing the model and the tokenizer.
9
  '''
10
+ model_name = 'facebook/incoder-6B'
11
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
12
+ model = AutoModelForCausalLM.from_pretrained(model_name, low_cpu_mem_usage=True)
13
  return [model, tokenizer]
14
 
15
  st.set_page_config(