Emalawi19 commited on
Commit
996b567
·
verified ·
1 Parent(s): 0eb71cb

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -12
app.py CHANGED
@@ -2,29 +2,52 @@ import gradio as gr
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
 
 
5
  model_name = "deepseek-ai/deepseek-coder-1.3b-base"
6
 
 
7
  tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 
8
  model = AutoModelForCausalLM.from_pretrained(
9
  model_name,
10
- torch_dtype=torch.float16,
11
- device_map="auto"
12
  )
13
 
 
 
 
 
14
  def generate_code(prompt):
15
- inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
16
- outputs = model.generate(
17
- **inputs,
18
- max_new_tokens=200,
19
- temperature=0.7
20
- )
21
- return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 
 
 
 
 
 
 
 
22
 
 
23
  iface = gr.Interface(
24
  fn=generate_code,
25
- inputs="text",
26
- outputs="text",
27
- title="DeepSeek Coder 1.3B"
 
 
 
 
28
  )
29
 
 
30
  iface.launch()
 
2
  from transformers import AutoTokenizer, AutoModelForCausalLM
3
  import torch
4
 
5
+ # Model name
6
  model_name = "deepseek-ai/deepseek-coder-1.3b-base"
7
 
8
+ # Load tokenizer
9
  tokenizer = AutoTokenizer.from_pretrained(model_name)
10
+
11
+ # Load model (optimized for CPU)
12
  model = AutoModelForCausalLM.from_pretrained(
13
  model_name,
14
+ torch_dtype=torch.float32, # CPU friendly
15
+ low_cpu_mem_usage=True
16
  )
17
 
18
+ model.to("cpu")
19
+ model.eval()
20
+
21
+ # Generate function
22
  def generate_code(prompt):
23
+ if not prompt.strip():
24
+ return "Please enter a prompt."
25
+
26
+ inputs = tokenizer(prompt, return_tensors="pt")
27
+
28
+ with torch.no_grad():
29
+ outputs = model.generate(
30
+ **inputs,
31
+ max_new_tokens=100, # lower = faster
32
+ temperature=0.5,
33
+ do_sample=True,
34
+ top_p=0.9
35
+ )
36
+
37
+ result = tokenizer.decode(outputs[0], skip_special_tokens=True)
38
+ return result
39
 
40
+ # Gradio UI
41
  iface = gr.Interface(
42
  fn=generate_code,
43
+ inputs=gr.Textbox(
44
+ lines=5,
45
+ placeholder="Write a coding prompt here..."
46
+ ),
47
+ outputs=gr.Textbox(lines=10),
48
+ title="DeepSeek Coder 1.3B (CPU Optimized)",
49
+ description="Generate code using DeepSeek Coder running on Hugging Face Spaces (CPU mode)."
50
  )
51
 
52
+ # Launch app
53
  iface.launch()