Pippoz commited on
Commit
c9ee852
1 Parent(s): 99b6634

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -9,11 +9,11 @@ import torch
9
  # model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype=torch.float16).cuda()
10
  # tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b", use_fast=False)
11
  # return model, tokenizer
12
-
 
13
  @st.cache(allow_output_mutation=True)
14
  def define_model(prompt):
15
- generator = pipeline('text-generation', model="facebook/opt-1.3b")
16
- answer = generator(prompt)
17
  return answer
18
 
19
 
 
9
  # model = AutoModelForCausalLM.from_pretrained("facebook/opt-1.3b", torch_dtype=torch.float16).cuda()
10
  # tokenizer = AutoTokenizer.from_pretrained("facebook/opt-1.3b", use_fast=False)
11
  # return model, tokenizer
12
+ generator = pipeline('text-generation', model="facebook/opt-1.3b", skip_special_tokens=True)
13
+
14
  @st.cache(allow_output_mutation=True)
15
  def define_model(prompt):
16
+ answer = generator(prompt, max_length=60)
 
17
  return answer
18
 
19