Apple commited on
Commit
9ba6e21
·
1 Parent(s): 2300419

Initial CADCoder Space with Gradio

Browse files
Files changed (2) hide show
  1. app.py +39 -0
  2. requirements.txt +5 -0
app.py ADDED
@@ -0,0 +1,39 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ import torch
4
+
5
+ # Load CAD-Coder model from Hugging Face
6
+ MODEL_NAME = "CADCODER/CAD-Coder"
7
+
8
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
9
+ model = AutoModelForCausalLM.from_pretrained(
10
+ MODEL_NAME,
11
+ trust_remote_code=True,
12
+ torch_dtype=torch.float16,
13
+ device_map="auto"
14
+ )
15
+
16
+ def generate_code(prompt):
17
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
18
+ with torch.no_grad():
19
+ outputs = model.generate(
20
+ **inputs,
21
+ max_new_tokens=512,
22
+ do_sample=True,
23
+ temperature=0.7,
24
+ top_p=0.9
25
+ )
26
+ return tokenizer.decode(outputs[0], skip_special_tokens=True)
27
+
28
+ # Gradio UI
29
+ demo = gr.Interface(
30
+ fn=generate_code,
31
+ inputs=gr.Textbox(lines=5, placeholder="Enter your CAD design prompt..."),
32
+ outputs="text",
33
+ title="CAD-Coder Inference",
34
+ description="Generate CAD code from natural language using CAD-Coder."
35
+ )
36
+
37
+ if __name__ == "__main__":
38
+ demo.launch()
39
+
requirements.txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ transformers
3
+ gradio
4
+ git+https://github.com/huggingface/accelerate.git
5
+