schuler commited on
Commit
3d905f8
·
verified ·
1 Parent(s): 72f7511

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -2,6 +2,7 @@ import gradio as gr
2
 
3
  import os, sys
4
  from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, pipeline
 
5
  import torch
6
  import spaces
7
  import psutil
@@ -13,7 +14,8 @@ REPO_NAME = 'schuler/experimental-JP47D20'
13
  # How to cache?
14
  @spaces.GPU()
15
  def load_model(repo_name):
16
- tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
 
17
  generator_conf = GenerationConfig.from_pretrained(repo_name)
18
  model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="eager")
19
  # model.to('cuda')
 
2
 
3
  import os, sys
4
  from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, pipeline
5
+ from transformers import LlamaTokenizer
6
  import torch
7
  import spaces
8
  import psutil
 
14
  # How to cache?
15
  @spaces.GPU()
16
  def load_model(repo_name):
17
+ # tokenizer = AutoTokenizer.from_pretrained(repo_name, trust_remote_code=True)
18
+ tokenizer = LlamaTokenizer.from_pretrained(repo_name, trust_remote_code=True)
19
  generator_conf = GenerationConfig.from_pretrained(repo_name)
20
  model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="eager")
21
  # model.to('cuda')