Update app.py

app.py

@@ -1,111 +1,41 @@
 import gradio as gr
-import torch
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 
-# Add specific import for Spaces GPU
-try:
-    from huggingface_hub.hf_api import HfFolder
-    from huggingface_hub import SpaceStage
-    print("HF中国镜像站 Hub imports successful")
-except ImportError:
-    print("Unable to import from huggingface_hub. Continuing without Spaces-specific handling.")
+# Define the model name
+model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"
 
-# Try the specific spaces import for GPU
-try:
-    from huggingface_hub import spaces
-    # If this works, we'll use it later
-    HAS_SPACES_GPU = True
-    print("Successfully imported spaces from huggingface_hub")
-except ImportError:
-    HAS_SPACES_GPU = False
-    print("spaces module not available from huggingface_hub")
+# Load tokenizer and model (disable fast tokenizer)
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
+model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
 
-# Define a function that handles the spaces.GPU decorator if available
-def main_wrapper():
-    # For Spaces that support the @spaces.GPU decorator
-    if HAS_SPACES_GPU:
-        print("Using @spaces.GPU decorator")
-        @spaces.GPU
-        def main_with_gpu():
-            run_app()
-        main_with_gpu()
-    else:
-        # Fallback for environments that don't support the decorator
-        print("Running without @spaces.GPU decorator")
-        run_app()
+# Create a text-generation pipeline
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-def run_app():
-    # Check GPU availability
-    if torch.cuda.is_available():
-        print(f"GPU is available: {torch.cuda.get_device_name(0)}")
-        device = "cuda"
-    else:
-        print("WARNING: GPU not available, using CPU instead. This will be very slow.")
-        device = "cpu"
-    
-    # Use a smaller model as a starting point
-    model_name = "ALLaM-AI/ALLaM-7B-Instruct-preview"  # Start with a smaller model to test
-    
-    try:
-        # Load tokenizer with appropriate options
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        
-        # Load model with appropriate device mapping
-        model_kwargs = {
-            "device_map": "auto" if device == "cuda" else None,
-        }
-        
-        if device == "cuda":
-            model_kwargs["torch_dtype"] = torch.float16
-        
-        model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
-        
-        # Create pipeline
-        text_generator = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            device=0 if device == "cuda" else -1
-        )
+def generate_story(prompt, max_length, temperature):
+    """
+    Generates a story based on the input prompt.
+    """
+    response = generator(
+        prompt,
+        max_length=max_length,
+        temperature=temperature,
+        do_sample=True
+    )
+    # Return the generated text
+    return response[0]["generated_text"]
 
-        # Define the generation function
-        def generate_story(prompt, max_length, temperature):
-            try:
-                response = text_generator(
-                    prompt,
-                    max_length=max_length,
-                    temperature=temperature,
-                    do_sample=True
-                )
-                return response[0]["generated_text"]
-            except Exception as e:
-                return f"Error generating text: {str(e)}"
-        
-        # Create the Gradio interface
-        demo = gr.Interface(
-            fn=generate_story,
-            inputs=[
-                gr.Textbox(lines=3, placeholder="Enter a story prompt here...", label="Story Prompt"),
-                gr.Slider(minimum=50, maximum=500, step=50, value=200, label="Max Length"),
-                gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Temperature")
-            ],
-            outputs=gr.Textbox(label="Generated Story"),
-            title="Story Generator",
-            description="Generate creative stories using a language model. Enter your prompt and let the AI craft a story for you!"
-        )
-        
-        # Launch the demo
-        demo.launch()
-        
-    except Exception as e:
-        print(f"Error initializing model: {str(e)}")
-        # Create a simple error interface
-        gr.Interface(
-            fn=lambda x: f"Error: {str(e)}. Please check the logs.",
-            inputs=gr.Textbox(),
-            outputs=gr.Textbox(),
-            title="Error Initializing Model"
-        ).launch()
+# Define the Gradio interface
+demo = gr.Interface(
+    fn=generate_story,
+    inputs=[
+        gr.Textbox(lines=3, placeholder="Enter a story prompt here...", label="Story Prompt"),
+        gr.Slider(minimum=50, maximum=1000, step=50, default=300, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=1.0, step=0.1, default=0.7, label="Temperature")
+    ],
+    outputs=gr.Textbox(label="Generated Story"),
+    title="Story Teller",
+    description="Generate creative stories using the ALLaM-7B model. Enter your prompt and let the model craft a story for you!"
+)
 
-if __name__ == "__main__":
-    main_wrapper()
+# Launch the app
+demo.launch()