AryanRathod3097 commited on
Commit
6f77195
·
verified ·
1 Parent(s): a58a992

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +125 -0
app.py ADDED
@@ -0,0 +1,125 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ codenyx.py – CodeNyx Coding Assistant
3
+ Powered by StarCoder2-15B-Instruct (fine-tuned on The Stack v2)
4
+ """
5
+ import gradio as gr
6
+ import torch
7
+ from transformers import (
8
+ AutoTokenizer,
9
+ AutoModelForCausalLM,
10
+ BitsAndBytesConfig,
11
+ TextIteratorStreamer
12
+ )
13
+ from threading import Thread
14
+
15
+ # ------------------------------------------------------------------
16
+ # 1. Branding & defaults
17
+ # ------------------------------------------------------------------
18
+ BOT_NAME = "CodeNyx"
19
+ SYSTEM_PROMPT = (f"You are {BOT_NAME}, an expert open-source coding assistant. "
20
+ "Always produce concise, runnable code snippets with explanations.")
21
+
22
+ MODEL_ID = "bigcode/starcoder2-15b-instruct-v0.1"
23
+ MAX_NEW_TOK = 2048
24
+ TEMPERATURE = 0.2
25
+ TOP_P = 0.9
26
+
27
+ # ------------------------------------------------------------------
28
+ # 2. Load model & tokenizer (4-bit quantized by default)
29
+ # ------------------------------------------------------------------
30
+ bnb_config = BitsAndBytesConfig(
31
+ load_in_4bit=True,
32
+ bnb_4bit_compute_dtype=torch.float16,
33
+ bnb_4bit_quant_type="nf4",
34
+ bnb_4bit_use_double_quant=True
35
+ )
36
+
37
+ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
38
+ if tokenizer.pad_token is None:
39
+ tokenizer.pad_token = tokenizer.eos_token
40
+
41
+ model = AutoModelForCausalLM.from_pretrained(
42
+ MODEL_ID,
43
+ quantization_config=bnb_config,
44
+ device_map="auto",
45
+ trust_remote_code=True
46
+ )
47
+
48
+ # ------------------------------------------------------------------
49
+ # 3. Chat helper
50
+ # ------------------------------------------------------------------
51
+ def build_prompt(history, user_input):
52
+ """Turn Gradio history into the model’s chat format."""
53
+ messages = [{"role": "system", "content": SYSTEM_PROMPT}]
54
+ for h in history:
55
+ messages.append({"role": "user", "content": h[0]})
56
+ messages.append({"role": "assistant", "content": h[1]})
57
+ messages.append({"role": "user", "content": user_input})
58
+ return tokenizer.apply_chat_template(
59
+ messages, tokenize=False, add_generation_prompt=True
60
+ )
61
+
62
+ def chat_fn(message, history):
63
+ prompt = build_prompt(history, message)
64
+ inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
65
+
66
+ streamer = TextIteratorStreamer(
67
+ tokenizer,
68
+ skip_prompt=True,
69
+ skip_special_tokens=True
70
+ )
71
+
72
+ generation_kwargs = dict(
73
+ **inputs,
74
+ max_new_tokens=MAX_NEW_TOK,
75
+ temperature=TEMPERATURE,
76
+ top_p=TOP_P,
77
+ do_sample=True,
78
+ pad_token_id=tokenizer.eos_token_id,
79
+ streamer=streamer
80
+ )
81
+ Thread(target=model.generate, kwargs=generation_kwargs).start()
82
+
83
+ partial = ""
84
+ for new_text in streamer:
85
+ partial += new_text
86
+ yield partial
87
+
88
+ # ------------------------------------------------------------------
89
+ # 4. Gradio UI
90
+ # ------------------------------------------------------------------
91
+ with gr.Blocks(title=f"{BOT_NAME} – AI Pair-Programmer") as demo:
92
+ gr.Markdown(f"""
93
+ # 🤖 {BOT_NAME}
94
+ *An open-source coding assistant trained on The Stack v2 – 15 B parameters, permissive licenses only.*
95
+ Type any programming question, and {BOT_NAME} will return **runnable code + concise explanations**.
96
+ """)
97
+
98
+ chatbot = gr.Chatbot(
99
+ label=f"{BOT_NAME} Chat",
100
+ height=600,
101
+ avatar_images=(None, "https://huggingface.co/spaces/bigcode/starcoder-logo/resolve/main/logo.png")
102
+ )
103
+
104
+ msg = gr.Textbox(
105
+ placeholder="Ask me to write, debug, or explain code...",
106
+ lines=2,
107
+ show_label=False,
108
+ container=False
109
+ )
110
+
111
+ clear = gr.Button("🗑️ Clear history")
112
+
113
+ def user(user_message, history):
114
+ return "", history + [[user_message, ""]]
115
+
116
+ msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
117
+ chat_fn, [msg, chatbot], chatbot
118
+ )
119
+ clear.click(lambda: None, None, chatbot, queue=False)
120
+
121
+ # ------------------------------------------------------------------
122
+ # 5. Launch
123
+ # ------------------------------------------------------------------
124
+ if __name__ == "__main__":
125
+ demo.queue().launch(server_name="0.0.0.0", server_port=7860, share=True)