We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
How to specify CPU or GPU?
I'd like to know how to move the model between CPU and GPU. Alternatively, I'd like to know if it's possible to specify the device before instantiating a model. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flair/models/language_model.py`
Content:
```
1 from pathlib import Path
2 
3 import torch.nn as nn
4 import torch
5 import math
6 from typing import Union, Tuple
7 from typing import List
8 
9 from torch.optim import Optimizer
10 
11 import flair
12 from flair.data import Dictionary
13 
14 
15 class LanguageModel(nn.Module):
16     """Container module with an encoder, a recurrent module, and a decoder."""
17 
18     def __init__(
19         self,
20         dictionary: Dictionary,
21         is_forward_lm: bool,
22         hidden_size: int,
23         nlayers: int,
24         embedding_size: int = 100,
25         nout=None,
26         dropout=0.1,
27     ):
28 
29         super(LanguageModel, self).__init__()
30 
31         self.dictionary = dictionary
32         self.is_forward_lm: bool = is_forward_lm
33 
34         self.dropout = dropout
35         self.hidden_size = hidden_size
36         self.embedding_size = embedding_size
37         self.nlayers = nlayers
38 
39         self.drop = nn.Dropout(dropout)
40         self.encoder = nn.Embedding(len(dictionary), embedding_size)
41 
42         if nlayers == 1:
43             self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers)
44         else:
45             self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers, dropout=dropout)
46 
47         self.hidden = None
48 
49         self.nout = nout
50         if nout is not None:
51             self.proj = nn.Linear(hidden_size, nout)
52             self.initialize(self.proj.weight)
53             self.decoder = nn.Linear(nout, len(dictionary))
54         else:
55             self.proj = None
56             self.decoder = nn.Linear(hidden_size, len(dictionary))
57 
58         self.init_weights()
59 
60         # auto-spawn on GPU if available
61         self.to(flair.device)
62 
63     def init_weights(self):
64         initrange = 0.1
65         self.encoder.weight.detach().uniform_(-initrange, initrange)
66         self.decoder.bias.detach().fill_(0)
67         self.decoder.weight.detach().uniform_(-initrange, initrange)
68 
69     def set_hidden(self, hidden):
70         self.hidden = hidden
71 
72     def forward(self, input, hidden, ordered_sequence_lengths=None):
73         encoded = self.encoder(input)
74         emb = self.drop(encoded)
75 
76         self.rnn.flatten_parameters()
77 
78         output, hidden = self.rnn(emb, hidden)
79 
80         if self.proj is not None:
81             output = self.proj(output)
82 
83         output = self.drop(output)
84 
85         decoded = self.decoder(
86             output.view(output.size(0) * output.size(1), output.size(2))
87         )
88 
89         return (
90             decoded.view(output.size(0), output.size(1), decoded.size(1)),
91             output,
92             hidden,
93         )
94 
95     def init_hidden(self, bsz):
96         weight = next(self.parameters()).detach()
97         return (
98             weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),
99             weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),
100         )
101 
102     def get_representation(self, strings: List[str], chars_per_chunk: int = 512):
103 
104         # cut up the input into chunks of max charlength = chunk_size
105         longest = len(strings[0])
106         chunks = []
107         splice_begin = 0
108         for splice_end in range(chars_per_chunk, longest, chars_per_chunk):
109             chunks.append([text[splice_begin:splice_end] for text in strings])
110             splice_begin = splice_end
111 
112         chunks.append([text[splice_begin:longest] for text in strings])
113         hidden = self.init_hidden(len(chunks[0]))
114 
115         output_parts = []
116 
117         # push each chunk through the RNN language model
118         for chunk in chunks:
119 
120             sequences_as_char_indices: List[List[int]] = []
121             for string in chunk:
122                 char_indices = [
123                     self.dictionary.get_idx_for_item(char) for char in string
124                 ]
125                 sequences_as_char_indices.append(char_indices)
126 
127             batch = torch.LongTensor(sequences_as_char_indices).transpose(0, 1)
128             batch = batch.to(flair.device)
129 
130             prediction, rnn_output, hidden = self.forward(batch, hidden)
131             rnn_output = rnn_output.detach()
132 
133             output_parts.append(rnn_output)
134 
135         # concatenate all chunks to make final output
136         output = torch.cat(output_parts)
137 
138         return output
139 
140     def get_output(self, text: str):
141         char_indices = [self.dictionary.get_idx_for_item(char) for char in text]
142         input_vector = torch.LongTensor([char_indices]).transpose(0, 1)
143 
144         hidden = self.init_hidden(1)
145         prediction, rnn_output, hidden = self.forward(input_vector, hidden)
146 
147         return self.repackage_hidden(hidden)
148 
149     def repackage_hidden(self, h):
150         """Wraps hidden states in new Variables, to detach them from their history."""
151         if type(h) == torch.Tensor:
152             return h.clone().detach()
153         else:
154             return tuple(self.repackage_hidden(v) for v in h)
155 
156     def initialize(self, matrix):
157         in_, out_ = matrix.size()
158         stdv = math.sqrt(3.0 / (in_ + out_))
159         matrix.detach().uniform_(-stdv, stdv)
160 
161     @classmethod
162     def load_language_model(cls, model_file: Union[Path, str]):
163 
164         state = torch.load(str(model_file), map_location=flair.device)
165 
166         model = LanguageModel(
167             state["dictionary"],
168             state["is_forward_lm"],
169             state["hidden_size"],
170             state["nlayers"],
171             state["embedding_size"],
172             state["nout"],
173             state["dropout"],
174         )
175         model.load_state_dict(state["state_dict"])
176         model.eval()
177         model.to(flair.device)
178 
179         return model
180 
181     @classmethod
182     def load_checkpoint(cls, model_file: Path):
183         state = torch.load(str(model_file), map_location=flair.device)
184 
185         epoch = state["epoch"] if "epoch" in state else None
186         split = state["split"] if "split" in state else None
187         loss = state["loss"] if "loss" in state else None
188         optimizer_state_dict = (
189             state["optimizer_state_dict"] if "optimizer_state_dict" in state else None
190         )
191 
192         model = LanguageModel(
193             state["dictionary"],
194             state["is_forward_lm"],
195             state["hidden_size"],
196             state["nlayers"],
197             state["embedding_size"],
198             state["nout"],
199             state["dropout"],
200         )
201         model.load_state_dict(state["state_dict"])
202         model.eval()
203         model.to(flair.device)
204 
205         return {
206             "model": model,
207             "epoch": epoch,
208             "split": split,
209             "loss": loss,
210             "optimizer_state_dict": optimizer_state_dict,
211         }
212 
213     def save_checkpoint(
214         self, file: Path, optimizer: Optimizer, epoch: int, split: int, loss: float
215     ):
216         model_state = {
217             "state_dict": self.state_dict(),
218             "dictionary": self.dictionary,
219             "is_forward_lm": self.is_forward_lm,
220             "hidden_size": self.hidden_size,
221             "nlayers": self.nlayers,
222             "embedding_size": self.embedding_size,
223             "nout": self.nout,
224             "dropout": self.dropout,
225             "optimizer_state_dict": optimizer.state_dict(),
226             "epoch": epoch,
227             "split": split,
228             "loss": loss,
229         }
230 
231         torch.save(model_state, str(file), pickle_protocol=4)
232 
233     def save(self, file: Path):
234         model_state = {
235             "state_dict": self.state_dict(),
236             "dictionary": self.dictionary,
237             "is_forward_lm": self.is_forward_lm,
238             "hidden_size": self.hidden_size,
239             "nlayers": self.nlayers,
240             "embedding_size": self.embedding_size,
241             "nout": self.nout,
242             "dropout": self.dropout,
243         }
244 
245         torch.save(model_state, str(file), pickle_protocol=4)
246 
247     def generate_text(
248         self,
249         prefix: str = "\n",
250         number_of_characters: int = 1000,
251         temperature: float = 1.0,
252         break_on_suffix=None,
253     ) -> Tuple[str, float]:
254 
255         if prefix == "":
256             prefix = "\n"
257 
258         with torch.no_grad():
259             characters = []
260 
261             idx2item = self.dictionary.idx2item
262 
263             # initial hidden state
264             hidden = self.init_hidden(1)
265 
266             if len(prefix) > 1:
267 
268                 char_tensors = []
269                 for character in prefix[:-1]:
270                     char_tensors.append(
271                         torch.tensor(self.dictionary.get_idx_for_item(character))
272                         .unsqueeze(0)
273                         .unsqueeze(0)
274                     )
275 
276                 input = torch.cat(char_tensors)
277                 if torch.cuda.is_available():
278                     input = input.cuda()
279 
280                 prediction, _, hidden = self.forward(input, hidden)
281 
282             input = (
283                 torch.tensor(self.dictionary.get_idx_for_item(prefix[-1]))
284                 .unsqueeze(0)
285                 .unsqueeze(0)
286             )
287 
288             log_prob = 0.0
289 
290             for i in range(number_of_characters):
291 
292                 if torch.cuda.is_available():
293                     input = input.cuda()
294 
295                 # get predicted weights
296                 prediction, _, hidden = self.forward(input, hidden)
297                 prediction = prediction.squeeze().detach()
298                 decoder_output = prediction
299 
300                 # divide by temperature
301                 prediction = prediction.div(temperature)
302 
303                 # to prevent overflow problem with small temperature values, substract largest value from all
304                 # this makes a vector in which the largest value is 0
305                 max = torch.max(prediction)
306                 prediction -= max
307 
308                 # compute word weights with exponential function
309                 word_weights = prediction.exp().cpu()
310 
311                 # try sampling multinomial distribution for next character
312                 try:
313                     word_idx = torch.multinomial(word_weights, 1)[0]
314                 except:
315                     word_idx = torch.tensor(0)
316 
317                 # print(word_idx)
318                 prob = decoder_output[word_idx]
319                 log_prob += prob
320 
321                 input = word_idx.detach().unsqueeze(0).unsqueeze(0)
322                 word = idx2item[word_idx].decode("UTF-8")
323                 characters.append(word)
324 
325                 if break_on_suffix is not None:
326                     if "".join(characters).endswith(break_on_suffix):
327                         break
328 
329             text = prefix + "".join(characters)
330 
331             log_prob = log_prob.item()
332             log_prob /= len(characters)
333 
334             if not self.is_forward_lm:
335                 text = text[::-1]
336 
337             return text, log_prob
338 
339     def calculate_perplexity(self, text: str) -> float:
340 
341         if not self.is_forward_lm:
342             text = text[::-1]
343 
344         # input ids
345         input = torch.tensor(
346             [self.dictionary.get_idx_for_item(char) for char in text[:-1]]
347         ).unsqueeze(1)
348         input = input.to(flair.device)
349 
350         # push list of character IDs through model
351         hidden = self.init_hidden(1)
352         prediction, _, hidden = self.forward(input, hidden)
353 
354         # the target is always the next character
355         targets = torch.tensor(
356             [self.dictionary.get_idx_for_item(char) for char in text[1:]]
357         )
358         targets = targets.to(flair.device)
359 
360         # use cross entropy loss to compare output of forward pass with targets
361         cross_entroy_loss = torch.nn.CrossEntropyLoss()
362         loss = cross_entroy_loss(
363             prediction.view(-1, len(self.dictionary)), targets
364         ).item()
365 
366         # exponentiate cross-entropy loss to calculate perplexity
367         perplexity = math.exp(loss)
368 
369         return perplexity
370
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flair/models/language_model.py b/flair/models/language_model.py
--- a/flair/models/language_model.py
+++ b/flair/models/language_model.py
@@ -273,9 +273,7 @@
                         .unsqueeze(0)
                     )
 
-                input = torch.cat(char_tensors)
-                if torch.cuda.is_available():
-                    input = input.cuda()
+                input = torch.cat(char_tensors).to(flair.device)
 
                 prediction, _, hidden = self.forward(input, hidden)
 
@@ -289,8 +287,7 @@
 
             for i in range(number_of_characters):
 
-                if torch.cuda.is_available():
-                    input = input.cuda()
+                input = input.to(flair.device)
 
                 # get predicted weights
                 prediction, _, hidden = self.forward(input, hidden)

{"golden_diff": "diff --git a/flair/models/language_model.py b/flair/models/language_model.py\n--- a/flair/models/language_model.py\n+++ b/flair/models/language_model.py\n@@ -273,9 +273,7 @@\n                         .unsqueeze(0)\n                     )\n \n-                input = torch.cat(char_tensors)\n-                if torch.cuda.is_available():\n-                    input = input.cuda()\n+                input = torch.cat(char_tensors).to(flair.device)\n \n                 prediction, _, hidden = self.forward(input, hidden)\n \n@@ -289,8 +287,7 @@\n \n             for i in range(number_of_characters):\n \n-                if torch.cuda.is_available():\n-                    input = input.cuda()\n+                input = input.to(flair.device)\n \n                 # get predicted weights\n                 prediction, _, hidden = self.forward(input, hidden)\n", "issue": "How to specify CPU or GPU?\nI'd like to know how to move the model between CPU and GPU. Alternatively, I'd like to know if it's possible to specify the device before instantiating a model. \n", "before_files": [{"content": "from pathlib import Path\n\nimport torch.nn as nn\nimport torch\nimport math\nfrom typing import Union, Tuple\nfrom typing import List\n\nfrom torch.optim import Optimizer\n\nimport flair\nfrom flair.data import Dictionary\n\n\nclass LanguageModel(nn.Module):\n    \"\"\"Container module with an encoder, a recurrent module, and a decoder.\"\"\"\n\n    def __init__(\n        self,\n        dictionary: Dictionary,\n        is_forward_lm: bool,\n        hidden_size: int,\n        nlayers: int,\n        embedding_size: int = 100,\n        nout=None,\n        dropout=0.1,\n    ):\n\n        super(LanguageModel, self).__init__()\n\n        self.dictionary = dictionary\n        self.is_forward_lm: bool = is_forward_lm\n\n        self.dropout = dropout\n        self.hidden_size = hidden_size\n        self.embedding_size = embedding_size\n        self.nlayers = nlayers\n\n        self.drop = nn.Dropout(dropout)\n        self.encoder = nn.Embedding(len(dictionary), embedding_size)\n\n        if nlayers == 1:\n            self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers)\n        else:\n            self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers, dropout=dropout)\n\n        self.hidden = None\n\n        self.nout = nout\n        if nout is not None:\n            self.proj = nn.Linear(hidden_size, nout)\n            self.initialize(self.proj.weight)\n            self.decoder = nn.Linear(nout, len(dictionary))\n        else:\n            self.proj = None\n            self.decoder = nn.Linear(hidden_size, len(dictionary))\n\n        self.init_weights()\n\n        # auto-spawn on GPU if available\n        self.to(flair.device)\n\n    def init_weights(self):\n        initrange = 0.1\n        self.encoder.weight.detach().uniform_(-initrange, initrange)\n        self.decoder.bias.detach().fill_(0)\n        self.decoder.weight.detach().uniform_(-initrange, initrange)\n\n    def set_hidden(self, hidden):\n        self.hidden = hidden\n\n    def forward(self, input, hidden, ordered_sequence_lengths=None):\n        encoded = self.encoder(input)\n        emb = self.drop(encoded)\n\n        self.rnn.flatten_parameters()\n\n        output, hidden = self.rnn(emb, hidden)\n\n        if self.proj is not None:\n            output = self.proj(output)\n\n        output = self.drop(output)\n\n        decoded = self.decoder(\n            output.view(output.size(0) * output.size(1), output.size(2))\n        )\n\n        return (\n            decoded.view(output.size(0), output.size(1), decoded.size(1)),\n            output,\n            hidden,\n        )\n\n    def init_hidden(self, bsz):\n        weight = next(self.parameters()).detach()\n        return (\n            weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),\n            weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),\n        )\n\n    def get_representation(self, strings: List[str], chars_per_chunk: int = 512):\n\n        # cut up the input into chunks of max charlength = chunk_size\n        longest = len(strings[0])\n        chunks = []\n        splice_begin = 0\n        for splice_end in range(chars_per_chunk, longest, chars_per_chunk):\n            chunks.append([text[splice_begin:splice_end] for text in strings])\n            splice_begin = splice_end\n\n        chunks.append([text[splice_begin:longest] for text in strings])\n        hidden = self.init_hidden(len(chunks[0]))\n\n        output_parts = []\n\n        # push each chunk through the RNN language model\n        for chunk in chunks:\n\n            sequences_as_char_indices: List[List[int]] = []\n            for string in chunk:\n                char_indices = [\n                    self.dictionary.get_idx_for_item(char) for char in string\n                ]\n                sequences_as_char_indices.append(char_indices)\n\n            batch = torch.LongTensor(sequences_as_char_indices).transpose(0, 1)\n            batch = batch.to(flair.device)\n\n            prediction, rnn_output, hidden = self.forward(batch, hidden)\n            rnn_output = rnn_output.detach()\n\n            output_parts.append(rnn_output)\n\n        # concatenate all chunks to make final output\n        output = torch.cat(output_parts)\n\n        return output\n\n    def get_output(self, text: str):\n        char_indices = [self.dictionary.get_idx_for_item(char) for char in text]\n        input_vector = torch.LongTensor([char_indices]).transpose(0, 1)\n\n        hidden = self.init_hidden(1)\n        prediction, rnn_output, hidden = self.forward(input_vector, hidden)\n\n        return self.repackage_hidden(hidden)\n\n    def repackage_hidden(self, h):\n        \"\"\"Wraps hidden states in new Variables, to detach them from their history.\"\"\"\n        if type(h) == torch.Tensor:\n            return h.clone().detach()\n        else:\n            return tuple(self.repackage_hidden(v) for v in h)\n\n    def initialize(self, matrix):\n        in_, out_ = matrix.size()\n        stdv = math.sqrt(3.0 / (in_ + out_))\n        matrix.detach().uniform_(-stdv, stdv)\n\n    @classmethod\n    def load_language_model(cls, model_file: Union[Path, str]):\n\n        state = torch.load(str(model_file), map_location=flair.device)\n\n        model = LanguageModel(\n            state[\"dictionary\"],\n            state[\"is_forward_lm\"],\n            state[\"hidden_size\"],\n            state[\"nlayers\"],\n            state[\"embedding_size\"],\n            state[\"nout\"],\n            state[\"dropout\"],\n        )\n        model.load_state_dict(state[\"state_dict\"])\n        model.eval()\n        model.to(flair.device)\n\n        return model\n\n    @classmethod\n    def load_checkpoint(cls, model_file: Path):\n        state = torch.load(str(model_file), map_location=flair.device)\n\n        epoch = state[\"epoch\"] if \"epoch\" in state else None\n        split = state[\"split\"] if \"split\" in state else None\n        loss = state[\"loss\"] if \"loss\" in state else None\n        optimizer_state_dict = (\n            state[\"optimizer_state_dict\"] if \"optimizer_state_dict\" in state else None\n        )\n\n        model = LanguageModel(\n            state[\"dictionary\"],\n            state[\"is_forward_lm\"],\n            state[\"hidden_size\"],\n            state[\"nlayers\"],\n            state[\"embedding_size\"],\n            state[\"nout\"],\n            state[\"dropout\"],\n        )\n        model.load_state_dict(state[\"state_dict\"])\n        model.eval()\n        model.to(flair.device)\n\n        return {\n            \"model\": model,\n            \"epoch\": epoch,\n            \"split\": split,\n            \"loss\": loss,\n            \"optimizer_state_dict\": optimizer_state_dict,\n        }\n\n    def save_checkpoint(\n        self, file: Path, optimizer: Optimizer, epoch: int, split: int, loss: float\n    ):\n        model_state = {\n            \"state_dict\": self.state_dict(),\n            \"dictionary\": self.dictionary,\n            \"is_forward_lm\": self.is_forward_lm,\n            \"hidden_size\": self.hidden_size,\n            \"nlayers\": self.nlayers,\n            \"embedding_size\": self.embedding_size,\n            \"nout\": self.nout,\n            \"dropout\": self.dropout,\n            \"optimizer_state_dict\": optimizer.state_dict(),\n            \"epoch\": epoch,\n            \"split\": split,\n            \"loss\": loss,\n        }\n\n        torch.save(model_state, str(file), pickle_protocol=4)\n\n    def save(self, file: Path):\n        model_state = {\n            \"state_dict\": self.state_dict(),\n            \"dictionary\": self.dictionary,\n            \"is_forward_lm\": self.is_forward_lm,\n            \"hidden_size\": self.hidden_size,\n            \"nlayers\": self.nlayers,\n            \"embedding_size\": self.embedding_size,\n            \"nout\": self.nout,\n            \"dropout\": self.dropout,\n        }\n\n        torch.save(model_state, str(file), pickle_protocol=4)\n\n    def generate_text(\n        self,\n        prefix: str = \"\\n\",\n        number_of_characters: int = 1000,\n        temperature: float = 1.0,\n        break_on_suffix=None,\n    ) -> Tuple[str, float]:\n\n        if prefix == \"\":\n            prefix = \"\\n\"\n\n        with torch.no_grad():\n            characters = []\n\n            idx2item = self.dictionary.idx2item\n\n            # initial hidden state\n            hidden = self.init_hidden(1)\n\n            if len(prefix) > 1:\n\n                char_tensors = []\n                for character in prefix[:-1]:\n                    char_tensors.append(\n                        torch.tensor(self.dictionary.get_idx_for_item(character))\n                        .unsqueeze(0)\n                        .unsqueeze(0)\n                    )\n\n                input = torch.cat(char_tensors)\n                if torch.cuda.is_available():\n                    input = input.cuda()\n\n                prediction, _, hidden = self.forward(input, hidden)\n\n            input = (\n                torch.tensor(self.dictionary.get_idx_for_item(prefix[-1]))\n                .unsqueeze(0)\n                .unsqueeze(0)\n            )\n\n            log_prob = 0.0\n\n            for i in range(number_of_characters):\n\n                if torch.cuda.is_available():\n                    input = input.cuda()\n\n                # get predicted weights\n                prediction, _, hidden = self.forward(input, hidden)\n                prediction = prediction.squeeze().detach()\n                decoder_output = prediction\n\n                # divide by temperature\n                prediction = prediction.div(temperature)\n\n                # to prevent overflow problem with small temperature values, substract largest value from all\n                # this makes a vector in which the largest value is 0\n                max = torch.max(prediction)\n                prediction -= max\n\n                # compute word weights with exponential function\n                word_weights = prediction.exp().cpu()\n\n                # try sampling multinomial distribution for next character\n                try:\n                    word_idx = torch.multinomial(word_weights, 1)[0]\n                except:\n                    word_idx = torch.tensor(0)\n\n                # print(word_idx)\n                prob = decoder_output[word_idx]\n                log_prob += prob\n\n                input = word_idx.detach().unsqueeze(0).unsqueeze(0)\n                word = idx2item[word_idx].decode(\"UTF-8\")\n                characters.append(word)\n\n                if break_on_suffix is not None:\n                    if \"\".join(characters).endswith(break_on_suffix):\n                        break\n\n            text = prefix + \"\".join(characters)\n\n            log_prob = log_prob.item()\n            log_prob /= len(characters)\n\n            if not self.is_forward_lm:\n                text = text[::-1]\n\n            return text, log_prob\n\n    def calculate_perplexity(self, text: str) -> float:\n\n        if not self.is_forward_lm:\n            text = text[::-1]\n\n        # input ids\n        input = torch.tensor(\n            [self.dictionary.get_idx_for_item(char) for char in text[:-1]]\n        ).unsqueeze(1)\n        input = input.to(flair.device)\n\n        # push list of character IDs through model\n        hidden = self.init_hidden(1)\n        prediction, _, hidden = self.forward(input, hidden)\n\n        # the target is always the next character\n        targets = torch.tensor(\n            [self.dictionary.get_idx_for_item(char) for char in text[1:]]\n        )\n        targets = targets.to(flair.device)\n\n        # use cross entropy loss to compare output of forward pass with targets\n        cross_entroy_loss = torch.nn.CrossEntropyLoss()\n        loss = cross_entroy_loss(\n            prediction.view(-1, len(self.dictionary)), targets\n        ).item()\n\n        # exponentiate cross-entropy loss to calculate perplexity\n        perplexity = math.exp(loss)\n\n        return perplexity\n", "path": "flair/models/language_model.py"}], "after_files": [{"content": "from pathlib import Path\n\nimport torch.nn as nn\nimport torch\nimport math\nfrom typing import Union, Tuple\nfrom typing import List\n\nfrom torch.optim import Optimizer\n\nimport flair\nfrom flair.data import Dictionary\n\n\nclass LanguageModel(nn.Module):\n    \"\"\"Container module with an encoder, a recurrent module, and a decoder.\"\"\"\n\n    def __init__(\n        self,\n        dictionary: Dictionary,\n        is_forward_lm: bool,\n        hidden_size: int,\n        nlayers: int,\n        embedding_size: int = 100,\n        nout=None,\n        dropout=0.1,\n    ):\n\n        super(LanguageModel, self).__init__()\n\n        self.dictionary = dictionary\n        self.is_forward_lm: bool = is_forward_lm\n\n        self.dropout = dropout\n        self.hidden_size = hidden_size\n        self.embedding_size = embedding_size\n        self.nlayers = nlayers\n\n        self.drop = nn.Dropout(dropout)\n        self.encoder = nn.Embedding(len(dictionary), embedding_size)\n\n        if nlayers == 1:\n            self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers)\n        else:\n            self.rnn = nn.LSTM(embedding_size, hidden_size, nlayers, dropout=dropout)\n\n        self.hidden = None\n\n        self.nout = nout\n        if nout is not None:\n            self.proj = nn.Linear(hidden_size, nout)\n            self.initialize(self.proj.weight)\n            self.decoder = nn.Linear(nout, len(dictionary))\n        else:\n            self.proj = None\n            self.decoder = nn.Linear(hidden_size, len(dictionary))\n\n        self.init_weights()\n\n        # auto-spawn on GPU if available\n        self.to(flair.device)\n\n    def init_weights(self):\n        initrange = 0.1\n        self.encoder.weight.detach().uniform_(-initrange, initrange)\n        self.decoder.bias.detach().fill_(0)\n        self.decoder.weight.detach().uniform_(-initrange, initrange)\n\n    def set_hidden(self, hidden):\n        self.hidden = hidden\n\n    def forward(self, input, hidden, ordered_sequence_lengths=None):\n        encoded = self.encoder(input)\n        emb = self.drop(encoded)\n\n        self.rnn.flatten_parameters()\n\n        output, hidden = self.rnn(emb, hidden)\n\n        if self.proj is not None:\n            output = self.proj(output)\n\n        output = self.drop(output)\n\n        decoded = self.decoder(\n            output.view(output.size(0) * output.size(1), output.size(2))\n        )\n\n        return (\n            decoded.view(output.size(0), output.size(1), decoded.size(1)),\n            output,\n            hidden,\n        )\n\n    def init_hidden(self, bsz):\n        weight = next(self.parameters()).detach()\n        return (\n            weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),\n            weight.new(self.nlayers, bsz, self.hidden_size).zero_().clone().detach(),\n        )\n\n    def get_representation(self, strings: List[str], chars_per_chunk: int = 512):\n\n        # cut up the input into chunks of max charlength = chunk_size\n        longest = len(strings[0])\n        chunks = []\n        splice_begin = 0\n        for splice_end in range(chars_per_chunk, longest, chars_per_chunk):\n            chunks.append([text[splice_begin:splice_end] for text in strings])\n            splice_begin = splice_end\n\n        chunks.append([text[splice_begin:longest] for text in strings])\n        hidden = self.init_hidden(len(chunks[0]))\n\n        output_parts = []\n\n        # push each chunk through the RNN language model\n        for chunk in chunks:\n\n            sequences_as_char_indices: List[List[int]] = []\n            for string in chunk:\n                char_indices = [\n                    self.dictionary.get_idx_for_item(char) for char in string\n                ]\n                sequences_as_char_indices.append(char_indices)\n\n            batch = torch.LongTensor(sequences_as_char_indices).transpose(0, 1)\n            batch = batch.to(flair.device)\n\n            prediction, rnn_output, hidden = self.forward(batch, hidden)\n            rnn_output = rnn_output.detach()\n\n            output_parts.append(rnn_output)\n\n        # concatenate all chunks to make final output\n        output = torch.cat(output_parts)\n\n        return output\n\n    def get_output(self, text: str):\n        char_indices = [self.dictionary.get_idx_for_item(char) for char in text]\n        input_vector = torch.LongTensor([char_indices]).transpose(0, 1)\n\n        hidden = self.init_hidden(1)\n        prediction, rnn_output, hidden = self.forward(input_vector, hidden)\n\n        return self.repackage_hidden(hidden)\n\n    def repackage_hidden(self, h):\n        \"\"\"Wraps hidden states in new Variables, to detach them from their history.\"\"\"\n        if type(h) == torch.Tensor:\n            return h.clone().detach()\n        else:\n            return tuple(self.repackage_hidden(v) for v in h)\n\n    def initialize(self, matrix):\n        in_, out_ = matrix.size()\n        stdv = math.sqrt(3.0 / (in_ + out_))\n        matrix.detach().uniform_(-stdv, stdv)\n\n    @classmethod\n    def load_language_model(cls, model_file: Union[Path, str]):\n\n        state = torch.load(str(model_file), map_location=flair.device)\n\n        model = LanguageModel(\n            state[\"dictionary\"],\n            state[\"is_forward_lm\"],\n            state[\"hidden_size\"],\n            state[\"nlayers\"],\n            state[\"embedding_size\"],\n            state[\"nout\"],\n            state[\"dropout\"],\n        )\n        model.load_state_dict(state[\"state_dict\"])\n        model.eval()\n        model.to(flair.device)\n\n        return model\n\n    @classmethod\n    def load_checkpoint(cls, model_file: Path):\n        state = torch.load(str(model_file), map_location=flair.device)\n\n        epoch = state[\"epoch\"] if \"epoch\" in state else None\n        split = state[\"split\"] if \"split\" in state else None\n        loss = state[\"loss\"] if \"loss\" in state else None\n        optimizer_state_dict = (\n            state[\"optimizer_state_dict\"] if \"optimizer_state_dict\" in state else None\n        )\n\n        model = LanguageModel(\n            state[\"dictionary\"],\n            state[\"is_forward_lm\"],\n            state[\"hidden_size\"],\n            state[\"nlayers\"],\n            state[\"embedding_size\"],\n            state[\"nout\"],\n            state[\"dropout\"],\n        )\n        model.load_state_dict(state[\"state_dict\"])\n        model.eval()\n        model.to(flair.device)\n\n        return {\n            \"model\": model,\n            \"epoch\": epoch,\n            \"split\": split,\n            \"loss\": loss,\n            \"optimizer_state_dict\": optimizer_state_dict,\n        }\n\n    def save_checkpoint(\n        self, file: Path, optimizer: Optimizer, epoch: int, split: int, loss: float\n    ):\n        model_state = {\n            \"state_dict\": self.state_dict(),\n            \"dictionary\": self.dictionary,\n            \"is_forward_lm\": self.is_forward_lm,\n            \"hidden_size\": self.hidden_size,\n            \"nlayers\": self.nlayers,\n            \"embedding_size\": self.embedding_size,\n            \"nout\": self.nout,\n            \"dropout\": self.dropout,\n            \"optimizer_state_dict\": optimizer.state_dict(),\n            \"epoch\": epoch,\n            \"split\": split,\n            \"loss\": loss,\n        }\n\n        torch.save(model_state, str(file), pickle_protocol=4)\n\n    def save(self, file: Path):\n        model_state = {\n            \"state_dict\": self.state_dict(),\n            \"dictionary\": self.dictionary,\n            \"is_forward_lm\": self.is_forward_lm,\n            \"hidden_size\": self.hidden_size,\n            \"nlayers\": self.nlayers,\n            \"embedding_size\": self.embedding_size,\n            \"nout\": self.nout,\n            \"dropout\": self.dropout,\n        }\n\n        torch.save(model_state, str(file), pickle_protocol=4)\n\n    def generate_text(\n        self,\n        prefix: str = \"\\n\",\n        number_of_characters: int = 1000,\n        temperature: float = 1.0,\n        break_on_suffix=None,\n    ) -> Tuple[str, float]:\n\n        if prefix == \"\":\n            prefix = \"\\n\"\n\n        with torch.no_grad():\n            characters = []\n\n            idx2item = self.dictionary.idx2item\n\n            # initial hidden state\n            hidden = self.init_hidden(1)\n\n            if len(prefix) > 1:\n\n                char_tensors = []\n                for character in prefix[:-1]:\n                    char_tensors.append(\n                        torch.tensor(self.dictionary.get_idx_for_item(character))\n                        .unsqueeze(0)\n                        .unsqueeze(0)\n                    )\n\n                input = torch.cat(char_tensors).to(flair.device)\n\n                prediction, _, hidden = self.forward(input, hidden)\n\n            input = (\n                torch.tensor(self.dictionary.get_idx_for_item(prefix[-1]))\n                .unsqueeze(0)\n                .unsqueeze(0)\n            )\n\n            log_prob = 0.0\n\n            for i in range(number_of_characters):\n\n                input = input.to(flair.device)\n\n                # get predicted weights\n                prediction, _, hidden = self.forward(input, hidden)\n                prediction = prediction.squeeze().detach()\n                decoder_output = prediction\n\n                # divide by temperature\n                prediction = prediction.div(temperature)\n\n                # to prevent overflow problem with small temperature values, substract largest value from all\n                # this makes a vector in which the largest value is 0\n                max = torch.max(prediction)\n                prediction -= max\n\n                # compute word weights with exponential function\n                word_weights = prediction.exp().cpu()\n\n                # try sampling multinomial distribution for next character\n                try:\n                    word_idx = torch.multinomial(word_weights, 1)[0]\n                except:\n                    word_idx = torch.tensor(0)\n\n                # print(word_idx)\n                prob = decoder_output[word_idx]\n                log_prob += prob\n\n                input = word_idx.detach().unsqueeze(0).unsqueeze(0)\n                word = idx2item[word_idx].decode(\"UTF-8\")\n                characters.append(word)\n\n                if break_on_suffix is not None:\n                    if \"\".join(characters).endswith(break_on_suffix):\n                        break\n\n            text = prefix + \"\".join(characters)\n\n            log_prob = log_prob.item()\n            log_prob /= len(characters)\n\n            if not self.is_forward_lm:\n                text = text[::-1]\n\n            return text, log_prob\n\n    def calculate_perplexity(self, text: str) -> float:\n\n        if not self.is_forward_lm:\n            text = text[::-1]\n\n        # input ids\n        input = torch.tensor(\n            [self.dictionary.get_idx_for_item(char) for char in text[:-1]]\n        ).unsqueeze(1)\n        input = input.to(flair.device)\n\n        # push list of character IDs through model\n        hidden = self.init_hidden(1)\n        prediction, _, hidden = self.forward(input, hidden)\n\n        # the target is always the next character\n        targets = torch.tensor(\n            [self.dictionary.get_idx_for_item(char) for char in text[1:]]\n        )\n        targets = targets.to(flair.device)\n\n        # use cross entropy loss to compare output of forward pass with targets\n        cross_entroy_loss = torch.nn.CrossEntropyLoss()\n        loss = cross_entroy_loss(\n            prediction.view(-1, len(self.dictionary)), targets\n        ).item()\n\n        # exponentiate cross-entropy loss to calculate perplexity\n        perplexity = math.exp(loss)\n\n        return perplexity\n", "path": "flair/models/language_model.py"}]}