Jeff Myers II commited on
Commit
7f1b674
·
1 Parent(s): 5fba035

Added all categories back to demo.

Browse files
Files changed (2) hide show
  1. Gemma.py +0 -22
  2. News.py +4 -4
Gemma.py CHANGED
@@ -1,4 +1,3 @@
1
- # from transformers import AutoTokenizer, Gemma3ForCausalLM
2
  from transformers import pipeline
3
  from huggingface_hub import login
4
  import spaces
@@ -15,31 +14,10 @@ class GemmaLLM:
15
 
16
  model_id = "google/gemma-3-4b-it"
17
 
18
- # self.tokenizer = AutoTokenizer.from_pretrained(model_id)
19
- # self.model = Gemma3ForCausalLM.from_pretrained(
20
- # model_id,
21
- # device_map="cuda" if torch.cuda.is_available() else "cpu",
22
- # torch_dtype=torch.float16,
23
- # ).eval()
24
-
25
  self.model = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16, device="cuda")
26
 
27
  @spaces.GPU
28
  def generate(self, message) -> str:
29
- # inputs = self.tokenizer.apply_chat_template(
30
- # message,
31
- # add_generation_prompt=True,
32
- # tokenize=True,
33
- # return_dict=True,
34
- # return_tensors="pt",
35
- # ).to(self.model.device)
36
-
37
- # input_length = inputs["input_ids"].shape[1]
38
-
39
- # with torch.inference_mode():
40
- # outputs = self.model.generate(**inputs, max_new_tokens=1024)[0][input_length:]
41
- # outputs = self.tokenizer.decode(outputs, skip_special_tokens=True)
42
-
43
  outputs = self.model(message, max_new_tokens=1024)[0]["generated_text"]
44
 
45
  return outputs
 
 
1
  from transformers import pipeline
2
  from huggingface_hub import login
3
  import spaces
 
14
 
15
  model_id = "google/gemma-3-4b-it"
16
 
 
 
 
 
 
 
 
17
  self.model = pipeline("text-generation", model=model_id, torch_dtype=torch.bfloat16, device="cuda")
18
 
19
  @spaces.GPU
20
  def generate(self, message) -> str:
 
 
 
 
 
 
 
 
 
 
 
 
 
 
21
  outputs = self.model(message, max_new_tokens=1024)[0]["generated_text"]
22
 
23
  return outputs
News.py CHANGED
@@ -8,10 +8,10 @@ class News:
8
  __EX_SOURCES__ = ["ABC News", "Bloomberg", "The Hill", "Fox Sports", "Google News", "Newsweek", "Politico"]
9
  __CATEGORIES__ = [
10
  "General",
11
- # "Business",
12
- # "Entertainment",
13
- # "Health",
14
- # "Science",
15
  "Technology"
16
  ]
17
 
 
8
  __EX_SOURCES__ = ["ABC News", "Bloomberg", "The Hill", "Fox Sports", "Google News", "Newsweek", "Politico"]
9
  __CATEGORIES__ = [
10
  "General",
11
+ "Business",
12
+ "Entertainment",
13
+ "Health",
14
+ "Science",
15
  "Technology"
16
  ]
17