Switch to coding-focused models and remove GPU requirements

README.md

@@ -8,9 +8,9 @@ app_port: 5001
 pinned: false
 ---
 
-# AI Text Generation Flask API
+# AI Coding Assistant Flask API
 
-A Flask web application that serves AI text generation models via a REST API.
+A Flask web application that serves coding-focused AI models via a REST API. Optimized for code generation, debugging, and programming assistance.
 
 ## Features
 

app.py

@@ -17,18 +17,18 @@ model = None
 tokenizer = None
 
 def load_model():
-    """Load the Llama model and tokenizer"""
+    """Load the model and tokenizer"""
     global model, tokenizer
     
     try:
-        logger.info("Loading Microsoft DialoGPT model (Llama alternative)...")
-        # Using an ungated model that works similarly
-        model_name = "microsoft/DialoGPT-large"
+        logger.info("Loading Salesforce CodeT5+ model (good for coding tasks)...")
+        # Using CodeT5+ which is good for coding and doesn't require GPU
+        model_name = "Salesforce/codet5p-770m"
         
-        # Alternative ungated models you can try:
-        # model_name = "distilbert/distilgpt2"  # Smaller, faster
-        # model_name = "gpt2-large"  # GPT-2 Large
-        # model_name = "EleutherAI/gpt-neo-1.3B"  # GPT-Neo
+        # Alternative good coding models (ungated):
+        # model_name = "microsoft/CodeBERT-base"
+        # model_name = "EleutherAI/gpt-neo-1.3B"  # Better general model
+        # model_name = "microsoft/DialoGPT-medium"  # Better conversation
         
         # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -37,36 +37,47 @@ def load_model():
         if tokenizer.pad_token is None:
             tokenizer.pad_token = tokenizer.eos_token
         
-        # Load model with optimizations
+        # Load model without quantization (CPU compatible)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            load_in_8bit=True,  # Use 8-bit quantization to reduce memory usage
+            torch_dtype=torch.float32,  # Use float32 for CPU compatibility
+            device_map=None,  # Let PyTorch handle device placement
             trust_remote_code=True
         )
         
-        logger.info("Model loaded successfully!")
+        logger.info("CodeT5+ model loaded successfully!")
         
     except Exception as e:
-        logger.error(f"Error loading model: {str(e)}")
-        # Fallback to a simpler model
+        logger.error(f"Error loading CodeT5+: {str(e)}")
+        # Fallback to GPT-Neo (much better than GPT-2)
         try:
-            logger.info("Falling back to GPT-2...")
-            model_name = "gpt2"
+            logger.info("Falling back to GPT-Neo 1.3B (better than GPT-2)...")
+            model_name = "EleutherAI/gpt-neo-1.3B"
             tokenizer = AutoTokenizer.from_pretrained(model_name)
             if tokenizer.pad_token is None:
                 tokenizer.pad_token = tokenizer.eos_token
             
             model = AutoModelForCausalLM.from_pretrained(
                 model_name,
-                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-                device_map="auto" if torch.cuda.is_available() else None
+                torch_dtype=torch.float32,
+                device_map=None,
+                trust_remote_code=True
             )
-            logger.info("Fallback model loaded successfully!")
+            logger.info("GPT-Neo model loaded successfully!")
         except Exception as fallback_error:
-            logger.error(f"Fallback also failed: {str(fallback_error)}")
-            raise e
+            logger.error(f"GPT-Neo also failed, using GPT-2: {str(fallback_error)}")
+            # Final fallback to GPT-2
+            model_name = "gpt2"
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float32,
+                device_map=None
+            )
+            logger.info("GPT-2 fallback loaded successfully!")
 
 def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9):
     """Generate response using the loaded model"""
@@ -76,8 +87,13 @@ def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9):
         raise ValueError("Model not loaded. Please ensure the model is properly initialized.")
     
     try:
-        # Simple prompt formatting (works for most models)
-        formatted_prompt = f"User: {prompt}\nAssistant:"
+        # Enhanced prompt formatting for coding tasks
+        if any(keyword in prompt.lower() for keyword in ['code', 'python', 'function', 'class', 'def ', 'import', 'javascript', 'html', 'css']):
+            # Coding-specific prompt format
+            formatted_prompt = f"# Task: {prompt}\n# Solution:\n"
+        else:
+            # General conversation format
+            formatted_prompt = f"Human: {prompt}\n\nAssistant: "
         
         # Tokenize the input
         inputs = tokenizer.encode(formatted_prompt, return_tensors="pt")
@@ -85,7 +101,7 @@ def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9):
         # Move to the same device as the model
         inputs = inputs.to(model.device)
         
-        # Generate response
+        # Generate response with better parameters for code
         with torch.no_grad():
             outputs = model.generate(
                 inputs,
@@ -95,15 +111,18 @@ def generate_response(prompt, max_length=512, temperature=0.7, top_p=0.9):
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
                 eos_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.1
+                repetition_penalty=1.1,
+                no_repeat_ngram_size=2  # Avoid repetitive code
             )
         
         # Decode the response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        # Extract only the assistant's response
-        if "Assistant:" in response:
-            response = response.split("Assistant:")[-1].strip()
+        # Extract only the generated response
+        if "# Solution:\n" in response:
+            response = response.split("# Solution:\n")[-1].strip()
+        elif "Assistant: " in response:
+            response = response.split("Assistant: ")[-1].strip()
         elif formatted_prompt in response:
             response = response.replace(formatted_prompt, "").strip()