Update app.py

app.py

@@ -2,7 +2,7 @@ import spaces
 import gradio as gr
 import torch
 import transformers
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import os
 
 title = """# Welcome to 🌟Tonic's✨StarCoder
@@ -14,28 +14,13 @@ default_system_prompt = """SYSTEM: You are an AI that code. Answer with code."""
 
 model_path = "bigcode/starcoder2-15b"
 
-
 hf_token = os.getenv("HF_TOKEN")
 if not hf_token:
     raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
-#   trust_remote_code=True,
-    token=hf_token,
-)
-
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-
-# import gradio as gr
-# from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
-
-# checkpoint = "bigcode/starcoder2-15b"
-# quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-# tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-# model = AutoModelForCausalLM.from_pretrained(checkpoint, quantization_config=quantization_config).to("cuda")
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained( model_path, quantization_config=quantization_config).to("cuda")
 @spaces.GPU
 def generate_text(prompt, temperature, max_length):
     inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
@@ -45,14 +30,12 @@ def generate_text(prompt, temperature, max_length):
 def gradio_app():
     with gr.Blocks() as demo:
         gr.Markdown(title)
-        output = gr.Code(label="Generated Code", lines=40)
-        with gr.Row():
-            generate_btn = gr.Button("Generate")
+        prompt = gr.Textbox(label="Enter your code prompt", placeholder="def print_hello_world():")
         with gr.Row():
             temperature = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature")
             max_length = gr.Slider(minimum=100, maximum=1024, step=10, value=100, label="Generate Length")
-        with gr.Row():
-            prompt = gr.Textbox(label="Enter your code prompt", placeholder="def print_hello_world():")
+        generate_btn = gr.Button("Try✨StarCoder")
+        output = gr.Code(label="Generated Code", lines=40)
 
         generate_btn.click(
             fn=generate_text,
@@ -60,4 +43,7 @@ def gradio_app():
             outputs=output
         )
 
-    demo.launch()
+    demo.launch()
+
+if __name__ == "__main__":
+    gradio_app()