Jeff Myers II commited on
Commit
c247815
·
1 Parent(s): 64fde65

Completed Prototype

Browse files
Files changed (2) hide show
  1. Gemma.py +4 -1
  2. requirements.txt +0 -2
Gemma.py CHANGED
@@ -1,4 +1,5 @@
1
  from transformers import AutoTokenizer, Gemma3ForCausalLM
 
2
  import torch
3
  import json
4
  import os
@@ -8,6 +9,8 @@ __export__ = ["GemmaLLM"]
8
  class GemmaLLM:
9
 
10
  def __init__(self):
 
 
11
  model_id = "google/gemma-3-1b-it"
12
 
13
  self.tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -15,7 +18,7 @@ class GemmaLLM:
15
  model_id,
16
  device_map="cuda" if torch.cuda.is_available() else "cpu",
17
  torch_dtype=torch.float16,
18
- token=os.environ.get("GEMMA_TOKEN"),
19
  ).eval()
20
 
21
  def generate(self, message) -> str:
 
1
  from transformers import AutoTokenizer, Gemma3ForCausalLM
2
+ from huggingface_hub import login
3
  import torch
4
  import json
5
  import os
 
9
  class GemmaLLM:
10
 
11
  def __init__(self):
12
+ login(token=os.environ.get("GEMMA_TOKEN"))
13
+
14
  model_id = "google/gemma-3-1b-it"
15
 
16
  self.tokenizer = AutoTokenizer.from_pretrained(model_id)
 
18
  model_id,
19
  device_map="cuda" if torch.cuda.is_available() else "cpu",
20
  torch_dtype=torch.float16,
21
+ # token=os.environ.get("GEMMA_TOKEN"),
22
  ).eval()
23
 
24
  def generate(self, message) -> str:
requirements.txt CHANGED
@@ -1,5 +1,3 @@
1
- Flask==3.1.0
2
- flask_cors==5.0.1
3
  newsapi_python==0.2.7
4
  newspaper3k==0.2.8
5
  torch==2.6.0
 
 
 
1
  newsapi_python==0.2.7
2
  newspaper3k==0.2.8
3
  torch==2.6.0