Update app.py

app.py

@@ -104,15 +104,12 @@ class TransformerDecoder(nn.Module):
         return output
 
     @classmethod
-    def from_pretrained(cls, model_path: str, device: str = 'cpu'):
-        """Load a pretrained model from a directory"""
+    def from_pretrained(cls, model_id: str, device: str = 'cpu'):
+        """Load a pretrained model from Hugging Face Hub"""
         try:
-            # Load config
-            config_path = os.path.join(model_path, "config.json")
-            if not os.path.exists(config_path):
-                raise FileNotFoundError(f"Config not found at {config_path}")
-                
-            with open(config_path) as f:
+            # Download config
+            config_file = hf_hub_download(repo_id=model_id, filename="config.json")
+            with open(config_file) as f:
                 config = json.load(f)
             
             # Create model instance
@@ -126,25 +123,23 @@ class TransformerDecoder(nn.Module):
                 dropout=config.get('dropout', 0.1)
             )
             
-            # Load weights
-            weights_path = os.path.join(model_path, "pytorch_model.bin")
-            if not os.path.exists(weights_path):
-                raise FileNotFoundError(f"Weights not found at {weights_path}")
-                
-            state_dict = torch.load(weights_path, map_location=device)
+            # Download and load weights
+            weights_file = hf_hub_download(repo_id=model_id, filename="pytorch_model.bin")
+            state_dict = torch.load(weights_file, map_location=device)
             model.load_state_dict(state_dict)
             
             return model.to(device)
             
         except Exception as e:
-            raise Exception(f"Error loading model from {model_path}: {str(e)}")
+            raise Exception(f"Error loading model from {model_id}: {str(e)}")
 
 def generate_text(prompt, max_length=100, temperature=0.7):
     try:
         # Load model and tokenizer from Hugging Face Hub
         model_id = "ninagala/shakespeare-model"
-        tokenizer_file = hf_hub_download(repo_id=model_id, filename="tokenizer.json")
         
+        # Download files from hub
+        tokenizer_file = hf_hub_download(repo_id=model_id, filename="tokenizer.json")
         model = TransformerDecoder.from_pretrained(model_id)
         tokenizer = Tokenizer.from_file(tokenizer_file)
         
@@ -153,6 +148,8 @@ def generate_text(prompt, max_length=100, temperature=0.7):
         tokens = tokenizer.encode(prompt).ids
         input_ids = torch.tensor(tokens).unsqueeze(0)
         
+        generated_tokens = []
+        
         with torch.no_grad():
             for _ in range(max_length):
                 outputs = model(input_ids)
@@ -161,7 +158,10 @@ def generate_text(prompt, max_length=100, temperature=0.7):
                 next_token = torch.multinomial(probs, num_samples=1)
                 input_ids = torch.cat([input_ids, next_token], dim=1)
                 
-                if next_token.item() == tokenizer.token_to_id("[EOS]"):
+                token_id = next_token.item()
+                generated_tokens.append(token_id)
+                
+                if token_id == tokenizer.token_to_id("[EOS]"):
                     break
         
         return tokenizer.decode(input_ids[0].tolist())