Locutusque commited on
Commit
e355c04
·
verified ·
1 Parent(s): 0667ea2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -51,7 +51,7 @@ The role you will play follows below.
51
  # Install flash-attn
52
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
53
  # Initialize the model pipeline
54
- generator = pipeline('text-generation', model='Locutusque/Thespis-Llama-3.1-8B', torch_dtype=torch.bfloat16, token=os.getenv("TOKEN"))
55
  @spaces.GPU
56
  def generate_text(prompt, system_prompt, temperature, top_p, top_k, repetition_penalty, max_length):
57
  """
 
51
  # Install flash-attn
52
  subprocess.run('pip install flash-attn --no-build-isolation', env={'FLASH_ATTENTION_SKIP_CUDA_BUILD': "TRUE"}, shell=True)
53
  # Initialize the model pipeline
54
+ generator = pipeline('text-generation', model='Locutusque/Thespis-Llama-3.1-8B', torch_dtype=torch.bfloat16, token=os.getenv("TOKEN"), load_in_8bit=True)
55
  @spaces.GPU
56
  def generate_text(prompt, system_prompt, temperature, top_p, top_k, repetition_penalty, max_length):
57
  """