Spaces:
Build error
Build error
Update app.py
Browse files
app.py
CHANGED
@@ -8,7 +8,7 @@ HF_TOKEN = os.environ.get("HF_TOKEN", None)
|
|
8 |
|
9 |
# Load the tokenizer and model
|
10 |
tokenizer = GemmaTokenizer.from_pretrained("google/codegemma-7b-it")
|
11 |
-
model = AutoModelForCausalLM.from_pretrained("google/codegemma-7b-it")
|
12 |
|
13 |
|
14 |
def codegemma(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
|
|
|
8 |
|
9 |
# Load the tokenizer and model
|
10 |
tokenizer = GemmaTokenizer.from_pretrained("google/codegemma-7b-it")
|
11 |
+
model = AutoModelForCausalLM.from_pretrained("google/codegemma-7b-it", device_map="auto") #.to("cuda:0")
|
12 |
|
13 |
|
14 |
def codegemma(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
|