rsepulvedat commited on
Commit
e8695e0
1 Parent(s): 1f780ee
Files changed (1) hide show
  1. app.py +3 -1
app.py CHANGED
@@ -1,8 +1,10 @@
1
  import gradio as gr
2
  from transformers import pipeline
3
  import torch
 
4
 
5
- pipe = pipeline("text-generation", model="gplsi/Aitana-6.3B", torch_dtype=torch.bfloat16, device_map="auto")
 
6
 
7
  def predict(input_text):
8
  generation = pipe(input_text, max_new_tokens=50, repetition_penalty=1.2, top_k=50, top_p=0.95, do_sample=True,
 
1
  import gradio as gr
2
  from transformers import pipeline
3
  import torch
4
+ import os
5
 
6
+ pipe = pipeline("text-generation", model="gplsi/Aitana-6.3B", torch_dtype=torch.bfloat16, device_map="auto",
7
+ token=os.environ['gplsi_models'])
8
 
9
  def predict(input_text):
10
  generation = pipe(input_text, max_new_tokens=50, repetition_penalty=1.2, top_k=50, top_p=0.95, do_sample=True,