Create app.py

app.py

@@ -0,0 +1,185 @@
+import spaces # If using Hugging Face Spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import gradio as gr
+import os
+
+# --- Environment and PyTorch Configurations (Kept from your original code) ---
+os.putenv('TORCH_LINALG_PREFER_CUSOLVER','1')
+os.putenv('PYTORCH_CUDA_ALLOC_CONF','max_split_size_mb:128') # Be mindful of this with a larger model
+os.environ["SAFETENSORS_FAST_GPU"] = "1"
+os.putenv('HF_HUB_ENABLE_HF_TRANSFER','1')
+
+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 # Setting to False can sometimes improve performance if exact reproducibility isn't critical
+torch.backends.cudnn.benchmark = True     # Setting to True can help if input sizes are consistent
+# torch.backends.cuda.preferred_blas_library="cublas" # These are generally defaults or fine
+# torch.backends.cuda.preferred_linalg_library="cusolver" # These are generally defaults or fine
+torch.set_float32_matmul_precision("highest") # Or "high" if "highest" causes issues
+
+# --- Model and Tokenizer Configuration ---
+# ** MODIFICATION 1: Update the model_name **
+model_name = "FelixChao/vicuna-33b-coder"
+
+# ** DOCUMENTATION: Model and Tokenizer Loading **
+# Load model and tokenizer.
+# `torch_dtype="auto"` will attempt to use the optimal dtype (e.g., bfloat16 if available).
+# `.to('cuda', torch.bfloat16)` explicitly moves the model to CUDA and casts to bfloat16.
+#   - Note: `FelixChao/vicuna-33b-coder` is a large model.
+#     Loading in bfloat16 requires substantial VRAM (~66GB+).
+#     If you encounter OOM errors, consider quantization:
+#     e.g., `load_in_8bit=True` or `load_in_4bit=True` (requires `bitsandbytes` library).
+# ** MODIFICATION 2: Removed `trust_remote_code=True` (typically not needed for Vicuna) **
+print(f"Loading model: {model_name}")
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="auto",
+    # device_map="auto", # device_map="auto" can be helpful for very large models to distribute layers if you have multiple GPUs or for offloading.
+                         # For single GPU, explicit .to('cuda') is fine.
+    # trust_remote_code=True # Removed: Generally not needed for Vicuna/Llama models
+).to('cuda', torch.bfloat16) # Explicitly using bfloat16 as in original code
+
+# ** MODIFICATION 3: Removed `trust_remote_code=True` for tokenizer **
+print(f"Loading tokenizer: {model_name}")
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    # trust_remote_code=True, # Removed: Generally not needed
+    use_fast=True
+)
+
+# ** DOCUMENTATION: Pad Token **
+# Vicuna/Llama models usually use EOS token as PAD token if PAD is not explicitly set.
+# This is often handled internally by the tokenizer or generation functions.
+# If you encounter issues related to padding, you might need to set it explicitly:
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+    print(f"Tokenizer `pad_token` was None, set to `eos_token`: {tokenizer.eos_token}")
+
+# Ensure the model's pad_token_id is also configured if necessary,
+# though `generate` usually handles this.
+model.config.pad_token_id = tokenizer.pad_token_id
+
+
+@spaces.GPU(required=True) # For Hugging Face Spaces GPU allocation
+def generate_code(prompt: str) -> str:
+    """
+    Generates code based on the given prompt using the loaded Vicuna model.
+
+    Args:
+        prompt: The user's input prompt for code generation.
+
+    Returns:
+        The generated code as a string.
+    """
+    # ** MODIFICATION 4: Update the system prompt for Vicuna **
+    # ** DOCUMENTATION: Chat Template Messages **
+    # The `messages` format is used by `tokenizer.apply_chat_template`.
+    # The system prompt provides context/instructions to the model.
+    # For Vicuna, a generic helpful assistant prompt is suitable.
+    messages = [
+        {"role": "system", "content": "You are a helpful and proficient coding assistant."},
+        {"role": "user", "content": prompt}
+    ]
+
+    # ** DOCUMENTATION: Applying Chat Template **
+    # `apply_chat_template` formats the input messages according to the model's specific template.
+    # `tokenize=False` returns a formatted string.
+    # `add_generation_prompt=True` ensures the template is set up for the model to start generating a response.
+    try:
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except Exception as e:
+        print(f"Error applying chat template: {e}")
+        # Fallback or simpler prompt structure if template is missing/problematic
+        # This is a basic fallback, actual Vicuna instruction format might be more specific
+        # e.g., "USER: {prompt}\nASSISTANT:"
+        # However, `apply_chat_template` is the preferred method.
+        # If this fails, the tokenizer for `FelixChao/vicuna-33b-coder` might be missing its template.
+        return f"Error: Could not apply chat template. The model's tokenizer might be misconfigured. ({e})"
+
+
+    # ** DOCUMENTATION: Tokenization **
+    # Tokenize the formatted text and move inputs to the model's device (GPU).
+    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    # ** DOCUMENTATION: Code Generation **
+    # Generate code using the model.
+    # `torch.no_grad()` disables gradient calculations, saving memory and computation during inference.
+    # `max_new_tokens`: Maximum number of new tokens to generate.
+    # `min_new_tokens`: Minimum number of new tokens to generate.
+    # `do_sample=True`: Enables sampling for more diverse outputs. If False, uses greedy decoding.
+    # `low_memory` is not a standard Hugging Face `generate` parameter.
+    #                 It might have been intended for a specific version or a custom argument.
+    #                 I've removed it as it might cause an error with standard generate.
+    #                 If you intended to use a specific memory optimization, that needs to be handled
+    #                 via other means (like quantization or model offloading).
+    with torch.no_grad():
+        generated_ids = model.generate(
+            **model_inputs,
+            max_new_tokens=1024,
+            min_new_tokens=256,
+            do_sample=True,
+            temperature=0.7, # Common sampling parameter, you can tune this
+            top_p=0.9,       # Common sampling parameter, you can tune this
+            pad_token_id=tokenizer.eos_token_id # Important for open-ended generation
+            # guidance_scale = 3.8, # Typically for Classifier Free Guidance, might not apply here or require specific setup
+        )
+
+    # ** DOCUMENTATION: Decoding Response **
+    # Remove the input tokens from the generated output to get only the response.
+    # `generated_ids[0]` because we process one prompt at a time.
+    # `model_inputs.input_ids[0]` is the input part.
+    response_ids = generated_ids[0][len(model_inputs.input_ids[0]):]
+
+    # Decode the generated token IDs back into a string.
+    # `skip_special_tokens=True` removes tokens like <s>, </s>, <unk>, etc.
+    response = tokenizer.decode(response_ids, skip_special_tokens=True)
+
+    return response.strip()
+
+# --- Gradio Interface (Kept mostly from your original code) ---
+with gr.Blocks(title="Vicuna 33B Coder") as demo: # Updated title
+    with gr.Tab("Code Chat"):
+        gr.Markdown("# Vicuna 33B Coder\nProvide a prompt to generate code.")
+        with gr.Row():
+            prompt = gr.Textbox( # Changed to Textbox for potentially longer prompts
+                label="Prompt",
+                show_label=True, # Changed to True for clarity
+                lines=3,         # Allow a few lines for the prompt
+                placeholder="Enter your coding prompt here...",
+                # container=False, # container is not a standard param for Textbox in this context
+            )
+        run_button = gr.Button("Generate Code", variant="primary")
+        with gr.Row():
+            result = gr.Code( # Using gr.Code for better code display
+                label="Generated Code",
+                show_label=True, # Changed to True for clarity
+                language="python", # Default language, can be auto-detected or set
+                lines=20,
+                # container=False,
+            )
+
+        # Gradio event listener
+        gr.on(
+            triggers=[
+                run_button.click,
+                prompt.submit # Allow submitting with Enter key in Textbox
+            ],
+            fn=generate_code,
+            inputs=[prompt],
+            outputs=[result],
+            # api_name="generate_code" # Uncomment if you want to expose this as an API endpoint
+        )
+
+if __name__ == "__main__":
+    # ** DOCUMENTATION: Launching Gradio **
+    # `share=False` by default, set to True if you want a public link (requires internet).
+    # `debug=True` can be helpful for development to see more detailed errors.
+    # `server_name="0.0.0.0"` to make it accessible on your local network.
+    demo.launch(share=False, debug=True)