Update pipeline.py

pipeline.py

@@ -22,23 +22,27 @@ class LSTMTextGenerator(nn.Module, PyTorchModelHubMixin):
 
 class PreTrainedPipeline():
     def __init__(self, path=""):
-        self.model = model = LSTMTextGenerator.from_pretrained("miittnnss/lstm-textgen-pets")
-        self.char_to_index = {char: index for index, char in enumerate(chars)}
- 
+        self.model = LSTMTextGenerator.from_pretrained("miittnnss/lstm-textgen-pets")
+        self.chars = ""
+        self.char_to_index = {char: index for index, char in enumerate(self.chars)}
+        self.index_to_char = {index: char for char, index in self.char_to_index.items()}
+        self.output_size = len(chars)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
     def __call__(self, inputs: str):
-        seed_numerical_data = [char_to_index[char] for char in inputs]
+        seed_numerical_data = [self.char_to_index[char] for char in inputs]
         with torch.no_grad():
-            input_sequence = torch.LongTensor([seed_numerical_data]).to(device)
-            hidden = model.init_hidden(1)
+            input_sequence = torch.LongTensor([seed_numerical_data]).to(self.device)
+            hidden = self.model.init_hidden(1)
 
             generated_text = inputs  # Initialize generated text with seed text
-            temperature = 0.7 # Temperature for temperature sampling
+            temperature = 0.7  # Temperature for temperature sampling
 
             for _ in range(500):
-                output, hidden = model(input_sequence, hidden)
+                output, hidden = self.model(input_sequence, hidden)
                 probabilities = nn.functional.softmax(output[-1, 0] / temperature, dim=0).cpu().numpy()
-                predicted_index = random.choices(range(output_size), weights=probabilities, k=1)[0]
-                generated_text += index_to_char[predicted_index]  # Append the generated character to the text
-                input_sequence = torch.LongTensor([[predicted_index]]).to(device)
+                predicted_index = random.choices(range(self.output_size), weights=probabilities, k=1)[0]
+                generated_text += self.index_to_char[predicted_index]  # Append the generated character to the text
+                input_sequence = torch.LongTensor([[predicted_index]]).to(self.device)
 
-        return output
+        return generated_text