ShishuTripathi commited on
Commit
0f5018a
1 Parent(s): c719137

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -4,7 +4,7 @@ import gradio as gr
4
  from peft import PeftModel, PeftConfig
5
  from transformers import AutoModelForCausalLM, AutoTokenizer ,pipeline, BitsAndBytesConfig
6
  config = PeftConfig.from_pretrained("ShishuTripathi/entity_coder")
7
- model = AutoModelForCausalLM.from_pretrained("ybelkada/falcon-7b-sharded-bf16",trust_remote_code=True,torch_dtype=torch.bfloat16,load_in_8bit=True,device_map={"":"cpu"},)
8
  model = PeftModel.from_pretrained(model,"ShishuTripathi/entity_coder")
9
  tokenizer = AutoTokenizer.from_pretrained("ShishuTripathi/entity_coder")
10
  generator = pipeline('text-generation' , model = model, tokenizer =tokenizer, max_length = 50)
 
4
  from peft import PeftModel, PeftConfig
5
  from transformers import AutoModelForCausalLM, AutoTokenizer ,pipeline, BitsAndBytesConfig
6
  config = PeftConfig.from_pretrained("ShishuTripathi/entity_coder")
7
+ model = AutoModelForCausalLM.from_pretrained("ybelkada/falcon-7b-sharded-bf16",trust_remote_code=True,torch_dtype=torch.bfloat16,load_in_8bit=True)
8
  model = PeftModel.from_pretrained(model,"ShishuTripathi/entity_coder")
9
  tokenizer = AutoTokenizer.from_pretrained("ShishuTripathi/entity_coder")
10
  generator = pipeline('text-generation' , model = model, tokenizer =tokenizer, max_length = 50)