Update app.py

app.py

@@ -2,77 +2,93 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gc
+import os
+
+# Set environment variables for cache management
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/.cache/huggingface'
+os.environ['HF_HOME'] = '/tmp/.cache/huggingface'
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
+
+def cleanup_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+    gc.collect()
 
-# Initialize model and tokenizer globally
 def load_model():
-    model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
+    cleanup_memory()
+    model_name = "Qwen/Qwen1.5-0.5B"  # Using smaller model
     try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            use_fast=False,
+            cache_dir="/tmp/.cache/huggingface"
+        )
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
             device_map="auto",
             trust_remote_code=True,
-            torch_dtype=torch.float16
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            offload_folder="/tmp/offload",
+            offload_state_dict=True,
+            cache_dir="/tmp/.cache/huggingface"
         )
         return model, tokenizer
     except Exception as e:
         print(f"Error loading model: {str(e)}")
         return None, None
 
+# Initialize model and tokenizer
+print("Loading model...")
 model, tokenizer = load_model()
+print("Model loaded successfully!")
 
 def generate_code(prompt):
     try:
-        # Clear CUDA cache
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
+        cleanup_memory()
         
         # Prepare the prompt
         messages = [
-            {"role": "system", "content": "You are an expert Python developer. Generate clean, efficient, and well-commented code based on the user's requirements."},
+            {"role": "system", "content": "You are an expert Python developer. Generate clean, efficient, and well-commented code based on the user's requirements. Only provide the code without any explanations."},
             {"role": "user", "content": f"Create a Python tool for the following requirement: {prompt}"}
         ]
         
-        # Format the messages using the chat template
-        text = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
-        
         # Generate the response
-        model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
-        generated_ids = model.generate(
-            **model_inputs,
-            max_new_tokens=1024,
+        inputs = tokenizer.encode(str(messages), return_tensors="pt", truncation=True, max_length=512)
+        inputs = inputs.to(model.device)
+        
+        outputs = model.generate(
+            inputs,
+            max_length=1024,
             temperature=0.7,
             top_p=0.9,
             repetition_penalty=1.1,
-            do_sample=True
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1
         )
         
-        # Decode the response
-        response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract only the code part from the response
+        # Extract code from response
         code_start = response.find("```python")
-        code_end = response.find("```", code_start + 8)
+        code_end = response.find("```", code_start + 8) if code_start != -1 else -1
         
         if code_start != -1 and code_end != -1:
             code = response[code_start + 8:code_end].strip()
         else:
-            code = response
-            
+            # If no code block found, try to extract any Python-like code
+            code = "\n".join([line for line in response.split("\n") 
+                            if line.strip() and not line.startswith(("#", "//", "/*"))])
+        
+        cleanup_memory()
         return code
         
     except Exception as e:
+        cleanup_memory()
         raise gr.Error(f"Code generation failed: {str(e)}")
-    finally:
-        # Clear memory
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-            gc.collect()
 
 # Create the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
@@ -80,7 +96,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
         """
         # 🛠️ AI Tool Builder by Syncmerce
         
-        Generate production-ready Python code for your tools using AI. Simply describe what you want to build!
+        Generate Python code for your tools using AI. Simply describe what you want to build!
         """
     )
     
@@ -99,17 +115,16 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
             code_output = gr.Code(
                 label="Generated Code",
                 language="python",
-                lines=25,
+                lines=20,
                 show_label=True
             )
     
-    # Add examples
+    # Add examples (keeping them simple to save memory)
     gr.Examples(
         examples=[
-            ["Create a PDF text extractor that can process multiple files and save the extracted text to a txt file"],
-            ["Build a web scraper that extracts product prices from Amazon and saves results to CSV"],
-            ["Create an image processing tool that can resize multiple images and add watermarks"],
-            ["Build a YouTube video downloader with progress bar and quality selection"],
+            ["Create a simple PDF text extractor"],
+            ["Build a basic web scraper for product prices"],
+            ["Create an image resizing tool"],
         ],
         inputs=prompt_input,
         outputs=code_output,
@@ -132,23 +147,22 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple")) as demo:
         api_name="clear"
     )
     
-    # Add footer
     gr.Markdown(
         """
         ### Tips for better results:
-        - Be specific about the features you want
-        - Mention input/output formats
-        - Specify any special requirements or libraries
-        
-        *Note: Generated code may need minor adjustments based on your specific use case.*
+        - Keep your requirements clear and concise
+        - Specify input/output formats if needed
+        - Mention any specific libraries you want to use
         """
     )
 
-# Launch the app
+# Launch the app with optimized settings
 demo.launch(
     share=True,
     enable_queue=True,
     show_error=True,
     server_name="0.0.0.0",
-    server_port=7860
+    server_port=7860,
+    cache_examples=True,
+    max_threads=4
 )