Update app.py

app.py

@@ -1,7 +1,145 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import os # For Hugging Face token
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Import spaces for ZeroGPU if you need to decorate specific functions
+# For models loaded via transformers and run on a device managed by ZeroGPU,
+# explicit @spaces.GPU might not always be needed directly on the inference function
+# if the entire space is on ZeroGPU hardware. However, for clarity or complex setups:
+# import spaces # Uncomment if using @spaces.GPU decorator
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+# --- Configuration ---
+HF_TOKEN = os.getenv("HF_TOKEN") # Recommended to store your Hugging Face token as a Space secret
+
+MODEL_OPTIONS = {
+    "Qwen1.5-1.8B-Chat": "Qwen/Qwen1.5-1.8B-Chat",
+    "Qwen2.5-Coder-3B": "Qwen/Qwen2.5-Coder-3B", # Example for a Qwen code model around 3B params
+}
+
+# --- Model Loading Cache ---
+# This dictionary will cache loaded models and tokenizers to avoid reloading on every call
+loaded_models = {}
+
+def get_model_and_tokenizer(model_name_key):
+    if model_name_key not in loaded_models:
+        model_id = MODEL_OPTIONS[model_name_key]
+        print(f"Loading model: {model_id}...")
+        try:
+            # Ensure you have accepted the terms of use for these models on Hugging Face Hub
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype="auto", # Let transformers decide the best dtype
+                device_map="auto",  # Automatically maps model to available device (GPU on ZeroGPU)
+                token=HF_TOKEN # Use token if model is private or requires it
+            )
+            tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
+            loaded_models[model_name_key] = (model, tokenizer)
+            print(f"Model {model_id} loaded successfully.")
+        except Exception as e:
+            print(f"Error loading model {model_id}: {e}")
+            # Fallback or error handling
+            if model_name_key in loaded_models: # Remove if partially loaded
+                del loaded_models[model_name_key]
+            raise gr.Error(f"Failed to load model {model_name_key}. Please check the model ID and your Hugging Face token permissions. Error: {e}")
+    return loaded_models[model_name_key]
+
+# --- Inference Function ---
+# If you need finer-grained control over GPU allocation for specific parts:
+# @spaces.GPU(duration=120) # Example: Request GPU for 120 seconds for this function
+def generate_response(prompt_text, model_choice, max_new_tokens=512, temperature=0.7, top_p=0.9):
+    if not prompt_text:
+        return "Please enter a prompt."
+    if not model_choice:
+        return "Please select a model."
+
+    try:
+        model, tokenizer = get_model_and_tokenizer(model_choice)
+    except Exception as e:
+        return str(e) # Display loading error to user
+
+    device = model.device # Get the device the model is on
+
+    if "Chat" in model_choice: # Apply chat template for chat models
+        messages = [
+            {"role": "system", "content": "You are a helpful assistant."},
+            {"role": "user", "content": prompt_text}
+        ]
+        try:
+            input_text = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+        except Exception as e: # Fallback if apply_chat_template has issues or is not applicable
+            print(f"Warning: Could not apply chat template for {model_choice}: {e}. Using prompt as is.")
+            input_text = prompt_text
+
+    else: # For code or non-chat models, use the prompt directly or adjust as needed
+        input_text = prompt_text
+
+    model_inputs = tokenizer([input_text], return_tensors="pt").to(device)
+
+    try:
+        generated_ids = model.generate(
+            model_inputs.input_ids,
+            max_new_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            do_sample=True # Necessary for temperature and top_p to have an effect
+        )
+        # For some models, the input prompt is included in the generated_ids.
+        # We need to decode only the newly generated tokens.
+        # This slicing can vary based on the model and tokenizer.
+        # A common approach is to slice based on the input_ids length:
+        response_ids = generated_ids[0][model_inputs.input_ids.shape[-1]:]
+        response_text = tokenizer.decode(response_ids, skip_special_tokens=True)
+
+    except Exception as e:
+        print(f"Error during generation with {model_choice}: {e}")
+        return f"Error generating response: {e}"
+
+    return response_text
+
+# --- Gradio Interface ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# LLM Coding & Math Experiment")
+    gr.Markdown("Query Qwen1.5-1.8B-Chat or Qwen Code models using ZeroGPU.")
+
+    with gr.Row():
+        model_dropdown = gr.Dropdown(
+            label="Select Model",
+            choices=list(MODEL_OPTIONS.keys()),
+            value=list(MODEL_OPTIONS.keys())[0] # Default to the first model
+        )
+    with gr.Row():
+        prompt_input = gr.Textbox(label="Enter your prompt:", lines=4, placeholder="e.g., Write a Python function to calculate factorial, or What is the capital of France?")
+    with gr.Row():
+        output_text = gr.Textbox(label="Model Response:", lines=8, interactive=False)
+
+    with gr.Row():
+        submit_button = gr.Button("Generate Response", variant="primary")
+
+    with gr.Accordion("Advanced Settings", open=False):
+        max_new_tokens_slider = gr.Slider(minimum=32, maximum=2048, value=512, step=32, label="Max New Tokens")
+        temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.05, label="Temperature")
+        top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P")
+
+
+    # Event listener for the button
+    submit_button.click(
+        fn=generate_response,
+        inputs=[prompt_input, model_dropdown, max_new_tokens_slider, temperature_slider, top_p_slider],
+        outputs=output_text,
+        api_name="generate" # Exposes an API endpoint
+    )
+
+    gr.Markdown("## Notes:")
+    gr.Markdown(
+        "- Ensure you have accepted the terms of use for the selected Qwen models on the Hugging Face Hub.\n"
+        "- Model loading can take some time, especially on the first run or when switching models.\n"
+        "- This Space runs on ZeroGPU, which means GPU resources are allocated dynamically."
+    )
+
+if __name__ == "__main__":
+    demo.launch()