Cyrano2 commited on
Commit
5dd4ced
·
verified ·
1 Parent(s): 0c9eb4b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -7
app.py CHANGED
@@ -1,15 +1,13 @@
1
- import gradio as gr
2
- from transformers import AutoModelForCausalLM, AutoTokenizer
3
  import torch
4
- from accelerate import init_empty_weights
5
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 
 
6
 
7
  # Charger le modèle
8
  model_name = "bigcode/starcoder2-15b-instruct-v0.1"
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
-
11
- # Initialisation conditionnelle
12
- device = "cuda" if torch.cuda.is_available() else "cpu"
13
  model = AutoModelForCausalLM.from_pretrained(
14
  model_name,
15
  torch_dtype=torch.float16 if device == "cuda" else torch.float32
@@ -21,7 +19,6 @@ def generate_text(prompt):
21
  outputs = model.generate(inputs["input_ids"], max_length=200)
22
  return tokenizer.decode(outputs[0], skip_special_tokens=True)
23
 
24
-
25
  # Interface utilisateur Gradio
26
  interface = gr.Interface(
27
  fn=generate_text,
 
 
 
1
  import torch
 
2
  from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import gradio as gr
4
+
5
+ # Vérifiez si un GPU est disponible avec ZeroGPU
6
+ device = "cuda" if torch.cuda.is_available() else "cpu"
7
 
8
  # Charger le modèle
9
  model_name = "bigcode/starcoder2-15b-instruct-v0.1"
10
  tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
 
11
  model = AutoModelForCausalLM.from_pretrained(
12
  model_name,
13
  torch_dtype=torch.float16 if device == "cuda" else torch.float32
 
19
  outputs = model.generate(inputs["input_ids"], max_length=200)
20
  return tokenizer.decode(outputs[0], skip_special_tokens=True)
21
 
 
22
  # Interface utilisateur Gradio
23
  interface = gr.Interface(
24
  fn=generate_text,