Update app.py

app.py

@@ -1,96 +1,22 @@
-import gradio as gr
 import torch
-from llama_cpp import Llama # New import for GGUF models
-
-# Define the model name for the GGUF model.
-# IMPORTANT: This assumes you will upload the 'Magistral-Small-2506_gguf' file
-# directly into the root directory of your Hugging Face Space.
-# You will need to download this file from its Hugging Face repository (e.g., from the "Files and versions" tab)
-# and upload it to your Space, naming it exactly as it appears here.
-GGUF_MODEL_FILE = "Magistral-Small-2506_gguf" # Adjust if your uploaded file name is different
-
-# Global variable for the Llama model instance
-llm = None
-
-# Function to load the Llama GGUF model
-def load_model():
-    global llm
-    if llm is None:
-        print(f"Loading GGUF model: {GGUF_MODEL_FILE}...")
-        try:
-            # Initialize the Llama model.
-            # `model_path` must point to the local file path of your GGUF model.
-            # `n_gpu_layers` can be set to a positive integer to offload layers to GPU if available.
-            # Set to 0 for CPU-only inference (recommended for simplicity on free Spaces tiers).
-            # `n_ctx` defines the context window size. Adjust as needed for your use case.
-            llm = Llama(model_path=GGUF_MODEL_FILE, n_gpu_layers=0, n_ctx=2048)
-            print("GGUF Model loaded successfully.")
-        except Exception as e:
-            print(f"Error loading GGUF model: {e}")
-            raise RuntimeError(f"Failed to load GGUF model: {e}. Please ensure '{GGUF_MODEL_FILE}' is correctly uploaded and accessible.")
-
-# Call this function once at the start of the script to load the model.
-load_model()
-
-# This is the core function that will be exposed as an API endpoint.
-# It takes a prompt and generation parameters, and returns generated text.
-def generate_text(prompt: str, max_new_tokens: int = 100, temperature: float = 0.7, top_k: int = 50) -> str:
-    # Basic input validation for the prompt.
-    if not prompt:
-        return "Please enter a prompt to generate text!"
-
-    if llm is None:
-        return "Model not loaded. Please check Space logs for errors."
-
-    try:
-        # Generate text using the Llama model's create_completion method.
-        # `prompt` is the input text.
-        # `max_tokens` controls the length of the generated output.
-        # `temperature` controls randomness (higher = more creative).
-        # `top_k` filters the sampling pool.
-        # `stop` can be used to define tokens where the generation should stop (e.g., ["\nUser:"]).
-        # `echo=False` ensures the prompt is not repeated in the output.
-        output = llm.create_completion(
-            prompt=prompt,
-            max_tokens=max_new_tokens,
-            temperature=temperature,
-            top_k=top_k,
-            stop=["\nUser:", "##"], # Example stop sequences
-            echo=False
-        )
-
-        # The generated text is typically found in the 'choices' list of the output dictionary.
-        generated_text = output['choices'][0]['text']
-        return generated_text
-
-    except Exception as e:
-        # Log any errors that occur during text generation for debugging.
-        print(f"Error during text generation: {e}")
-        # Return an informative error message to the user/caller.
-        return f"An error occurred: {e}. Please try again with a different prompt or check the Space logs."
-
-# Create the Gradio interface.
-# This interface will automatically generate a web UI and an API endpoint.
-demo = gr.Interface(
-    fn=generate_text, # The Python function to expose.
-    inputs=[
-        # Input component for the prompt.
-        gr.Textbox(label="Enter your prompt here", lines=3),
-        # Slider for maximum number of new tokens to generate.
-        gr.Slider(minimum=10, maximum=500, value=100, label="Max New Tokens"),
-        # Slider for generation temperature (randomness).
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Temperature"),
-        # Slider for Top-K sampling (diversity).
-        gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K")
-    ],
-    outputs=gr.Textbox(label="Generated Text", lines=5), # Output component for generated text.
-    title="Magistral-Small-2506_gguf Text Generation API on Hugging Face Space",
-    description="Enter a prompt and Magistral-Small-2506_gguf will generate a response. Adjust parameters for different results. This function is also exposed as an API endpoint.",
-    allow_flagging="never" # Disables Gradio's data flagging feature.
-)
-
-# Launch the Gradio application.
-# `server_name="0.0.0.0"` is essential for Hugging Face Spaces to expose the app publicly.
-# `server_port=7860` is the default port used by Hugging Face Spaces.
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+from transformers import pipeline
+
+# Check for GPU
+if torch.cuda.is_available():
+    print(f"CUDA is available! Using {torch.cuda.get_device_name(0)}")
+    device = 0 # Use GPU
+else:
+    print("CUDA not available, using CPU.")
+    device = -1 # Use CPU
+
+# Load a text generation pipeline
+# For a free tier/small GPU, consider a smaller model like 'distilgpt2' or 'gpt2'
+# For larger GPUs, you can try models like 'meta-llama/Llama-2-7b-hf' (requires auth)
+# or 'mistralai/Mistral-7B-Instruct-v0.2'
+generator = pipeline('text-generation', model='distilgpt2', device=device) # or specify a larger model
+
+# Generate text
+prompt = "The quick brown fox jumps over the lazy dog because"
+result = generator(prompt, max_length=50, num_return_sequences=1)
+
+print(result[0]['generated_text'])