Jeff Myers II commited on
Commit
eebf9a7
·
1 Parent(s): 7e0c20d

Completed Prototype

Browse files
Files changed (2) hide show
  1. Gemma.py +2 -0
  2. requirements.txt +2 -1
Gemma.py CHANGED
@@ -1,5 +1,6 @@
1
  from transformers import AutoTokenizer, Gemma3ForCausalLM
2
  from huggingface_hub import login
 
3
  import torch
4
  import json
5
  import os
@@ -21,6 +22,7 @@ class GemmaLLM:
21
  # token=os.environ.get("GEMMA_TOKEN"),
22
  ).eval()
23
 
 
24
  def generate(self, message) -> str:
25
  inputs = self.tokenizer.apply_chat_template(
26
  message,
 
1
  from transformers import AutoTokenizer, Gemma3ForCausalLM
2
  from huggingface_hub import login
3
+ import spaces
4
  import torch
5
  import json
6
  import os
 
22
  # token=os.environ.get("GEMMA_TOKEN"),
23
  ).eval()
24
 
25
+ @spaces.GPU
26
  def generate(self, message) -> str:
27
  inputs = self.tokenizer.apply_chat_template(
28
  message,
requirements.txt CHANGED
@@ -4,4 +4,5 @@ newsapi_python==0.2.7
4
  newspaper3k==0.2.8
5
  transformers==4.50.0
6
  lxml_html_clean==0.4.1
7
- accelerate==1.5.2
 
 
4
  newspaper3k==0.2.8
5
  transformers==4.50.0
6
  lxml_html_clean==0.4.1
7
+ accelerate==1.5.2
8
+ spaces