Spaces:
Running
Running
Jeff Myers II
commited on
Commit
·
7f1b674
1
Parent(s):
5fba035
Added all categories back to demo.
Browse files
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 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
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 |
|