Create app.py

app.py

@@ -0,0 +1,164 @@
+import spaces # If using Hugging Face Spaces
+
+import os
+
+os.putenv('PYTORCH_NVML_BASED_CUDA_CHECK','1')
+os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
+alloc_conf_parts = [
+    'expandable_segments:True',
+    'pinned_use_background_threads:True'  # Specific to pinned memory.
+]
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(alloc_conf_parts)
+os.environ["SAFETENSORS_FAST_GPU"] = "1"
+os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig # Import BitsAndBytesConfig
+import torch
+import gradio as gr
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
+torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
+torch.backends.cudnn.allow_tf32 = False
+torch.backends.cudnn.deterministic = False
+torch.backends.cudnn.benchmark = True
+torch.set_float32_matmul_precision("highest")
+
+# --- Model and Tokenizer Configuration ---
+model_name = "FelixChao/vicuna-33b-coder"
+
+# --- Quantization Configuration (Example: 4-bit) ---
+# This section is included based on our previous discussion.
+# Remove or comment out if you are not using quantization.
+print("Setting up 4-bit quantization config...")
+quantization_config_4bit = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+
+print(f"Loading model: {model_name} with quantization")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quantization_config_4bit, # Comment out if not using quantization
+    device_map="auto",
+)
+
+print(f"Loading tokenizer: {model_name}")
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    use_fast=True
+)
+
+# ** MODIFICATION: Define and set the Vicuna chat template **
+# ** DOCUMENTATION: Chat Template **
+# Vicuna models expect a specific chat format. If the tokenizer doesn't have one
+# built-in, we need to set it manually.
+# This template handles a system prompt, user messages, and assistant responses.
+# It will also add the "ASSISTANT:" prompt for generation if needed.
+VICUNA_CHAT_TEMPLATE = (
+    "{% if messages[0]['role'] == 'system' %}"  # Check if the first message is a system prompt
+        "{{ messages[0]['content'] + '\\n\\n' }}"  # Add system prompt with two newlines
+        "{% set loop_messages = messages[1:] %}"  # Slice to loop over remaining messages
+    "{% else %}"
+        "{% set loop_messages = messages %}"  # No system prompt, loop over all messages
+    "{% endif %}"
+    "{% for message in loop_messages %}"  # Loop through user and assistant messages
+        "{% if message['role'] == 'user' %}"
+            "{{ 'USER: ' + message['content'].strip() + '\\n' }}"
+        "{% elif message['role'] == 'assistant' %}"
+            "{{ 'ASSISTANT: ' + message['content'].strip() + eos_token + '\\n' }}"
+        "{% endif %}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}"  # If we need to prompt the model for a response
+        "{% if messages[-1]['role'] != 'assistant' %}" # And the last message wasn't from the assistant
+            "{{ 'ASSISTANT:' }}"  # Add the assistant prompt
+        "{% endif %}"
+    "{% endif %}"
+)
+tokenizer.chat_template = VICUNA_CHAT_TEMPLATE
+print("Manually set Vicuna chat template on the tokenizer.")
+
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    # Also update the model config's pad_token_id if you are setting tokenizer.pad_token
+    # This is crucial if the model's config doesn't get updated automatically.
+    if model.config.pad_token_id is None:
+         model.config.pad_token_id = tokenizer.pad_token_id
+    print(f"Tokenizer `pad_token` was None, set to `eos_token`: {tokenizer.eos_token}")
+
+
+@spaces.GPU(required=True)
+def generate_code(prompt: str) -> str:
+    messages = [
+        {"role": "system", "content": "You are a helpful and proficient coding assistant."},
+        {"role": "user", "content": prompt}
+    ]
+    try:
+        # ** DOCUMENTATION: Applying Chat Template **
+        # Now that tokenizer.chat_template is set, this should work.
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True # Important to append "ASSISTANT:"
+        )
+        print(f"Formatted prompt using chat template:\n{text}") # For debugging
+    except Exception as e:
+        print(f"Error applying chat template: {e}")
+        # Provide a more informative error or fallback if needed
+        return f"Error: Could not apply chat template. Details: {e}. Ensure the tokenizer has a valid `chat_template` attribute."
+
+    # Determine device for inputs if model is on multiple devices
+    # For device_map="auto", input tensors should go to the device of the first model block.
+    input_device = model.hf_device_map.get("", next(iter(model.hf_device_map.values()))) if hasattr(model, "hf_device_map") else model.device
+
+    model_inputs = tokenizer([text], return_tensors="pt").to(input_device)
+
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **model_inputs, # Pass tokenized inputs
+            max_new_tokens=1024,
+            min_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id # Use EOS token for padding
+        )
+
+    response_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
+    response = tokenizer.decode(response_ids, skip_special_tokens=True)
+    return response.strip()
+
+# --- Gradio Interface ---
+with gr.Blocks(title="Vicuna 33B Coder") as demo:
+    with gr.Tab("Code Chat"):
+        gr.Markdown("# Vicuna 33B Coder\nProvide a prompt to generate code.")
+        with gr.Row():
+            prompt_input = gr.Textbox( # Renamed to avoid conflict with 'prompt' variable in function scope
+                label="Prompt",
+                show_label=True,
+                lines=3,
+                placeholder="Enter your coding prompt here...",
+            )
+        run_button = gr.Button("Generate Code", variant="primary")
+        with gr.Row():
+            result_output = gr.Code( # Renamed
+                label="Generated Code",
+                show_label=True,
+                language="python",
+                lines=20,
+            )
+        gr.on(
+            triggers=[
+                run_button.click,
+                prompt_input.submit
+            ],
+            fn=generate_code,
+            inputs=[prompt_input],
+            outputs=[result_output],
+        )
+
+if __name__ == "__main__":
+    demo.launch(share=False, debug=True)