Update app.py

app.py

@@ -5,6 +5,7 @@ from torch import nn
 import requests
 from pathlib import Path
 import logging
+import os
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -22,13 +23,28 @@ class TextGenerator(nn.Module):
         lstm_out, _ = self.lstm(x)
         return self.fc(lstm_out)
 
+def fix_state_dict(state_dict):
+    """Fix the state dict keys by removing any 'module.' prefix"""
+    new_state_dict = {}
+    for k, v in state_dict.items():
+        name = k.replace('module.', '')  # Remove 'module.' prefix if it exists
+        new_state_dict[name] = v
+    return new_state_dict
+
 def download_file(url, local_path):
     try:
-        response = requests.get(url)
-        response.raise_for_status()  # Raise an exception for bad status codes
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        total_size = int(response.headers.get('content-length', 0))
+        block_size = 8192
+        
         Path(local_path).parent.mkdir(parents=True, exist_ok=True)
+        
         with open(local_path, 'wb') as f:
-            f.write(response.content)
+            for data in response.iter_content(block_size):
+                f.write(data)
+                
         logger.info(f"Successfully downloaded {url} to {local_path}")
     except Exception as e:
         logger.error(f"Error downloading {url}: {str(e)}")
@@ -41,9 +57,9 @@ def load_model_and_tokenizers():
     
     # Default configuration values
     default_config = {
-        'vocab_size': 10000,  # Default vocabulary size
-        'embedding_dim': 256,  # Default embedding dimension
-        'hidden_dim': 512     # Default hidden dimension
+        'vocab_size': 10000,
+        'embedding_dim': 256,
+        'hidden_dim': 512
     }
     
     # URLs for the files
@@ -64,13 +80,19 @@ def load_model_and_tokenizers():
     
     try:
         # Load configuration
-        with open(cache_dir / "model_config.json", "r") as f:
-            config = json.load(f)
-            # Merge with default config
-            for key in default_config:
-                if key not in config:
-                    logger.warning(f"Configuration parameter '{key}' not found, using default value: {default_config[key]}")
-                    config[key] = default_config[key]
+        config_path = cache_dir / "model_config.json"
+        if config_path.exists():
+            with open(config_path, "r") as f:
+                config = json.load(f)
+        else:
+            logger.warning("Config file not found, using default configuration.")
+            config = default_config
+
+        # Merge with default config
+        for key in default_config:
+            if key not in config:
+                logger.warning(f"Configuration parameter '{key}' not found, using default value: {default_config[key]}")
+                config[key] = default_config[key]
     except Exception as e:
         logger.warning(f"Error loading config file: {str(e)}. Using default configuration.")
         config = default_config
@@ -97,10 +119,34 @@ def load_model_and_tokenizers():
             hidden_dim=config['hidden_dim']
         )
         
-        # Load model weights
-        model.load_state_dict(torch.load(cache_dir / "model.pt", map_location=torch.device('cpu')))
-        model.eval()
+        # Load model weights with proper error handling
+        model_path = cache_dir / "model.pt"
+        if not model_path.exists():
+            raise FileNotFoundError(f"Model file not found at {model_path}")
+            
+        # Try different loading approaches
+        try:
+            # Try loading as a complete model
+            loaded_model = torch.load(model_path, map_location=torch.device('cpu'))
+            if isinstance(loaded_model, dict):
+                # If it's a state dict
+                state_dict = fix_state_dict(loaded_model)
+                model.load_state_dict(state_dict)
+            else:
+                # If it's a complete model
+                model = loaded_model
+        except Exception as e:
+            logger.warning(f"First loading attempt failed: {str(e)}")
+            try:
+                # Try loading as a state dict directly
+                state_dict = torch.load(model_path, map_location=torch.device('cpu'))
+                state_dict = fix_state_dict(state_dict)
+                model.load_state_dict(state_dict)
+            except Exception as e2:
+                logger.error(f"Both loading attempts failed. Last error: {str(e2)}")
+                raise
         
+        model.eval()
         return model, word_to_int, int_to_word
         
     except Exception as e:
@@ -121,7 +167,7 @@ def generate_text(prompt, max_length=100):
         
         with torch.no_grad():
             for _ in range(max_length):
-                current_input = torch.tensor([generated_ids[-50:]])  # Use last 50 tokens as context
+                current_input = torch.tensor([generated_ids[-50:]])
                 outputs = model(current_input)
                 next_token_id = outputs[0, -1, :].argmax().item()
                 generated_ids.append(next_token_id)
@@ -135,7 +181,7 @@ def generate_text(prompt, max_length=100):
         
     except Exception as e:
         logger.error(f"Error in text generation: {str(e)}")
-        return f"Error generating text: {str(e)}"
+        return f"Error generating text: {str(e)}\nPlease check the logs for more details."
 
 # Create Gradio interface
 iface = gr.Interface(