fix torch typehint on some neurons BT-1329 (#1460)

opentensor · Jul 26, 2023 · a9954b6 · a9954b6
1 parent 56cda18
commit a9954b6
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 6 deletions.
diff --git a/neurons/text/prompting/miners/AI21/neuron.py b/neurons/text/prompting/miners/AI21/neuron.py
@@ -17,6 +17,7 @@
 
 import argparse
 import bittensor
+from torch import FloatTensor
 
 from typing import List, Dict
 from langchain.llms import AI21
@@ -57,7 +58,7 @@ def __init__(self):
         bittensor.logging.info("Model loaded!")
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass
 

diff --git a/neurons/text/prompting/miners/AlephAlpha/neuron.py b/neurons/text/prompting/miners/AlephAlpha/neuron.py
@@ -19,6 +19,7 @@
 import bittensor
 from rich import print
 from typing import List, Dict
+from torch import FloatTensor
 
 from langchain.llms import AlephAlpha
 
@@ -87,7 +88,7 @@ def __init__(self):
         )
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass
 

diff --git a/neurons/text/prompting/miners/cerebras/neuron.py b/neurons/text/prompting/miners/cerebras/neuron.py
@@ -20,6 +20,7 @@
 import bittensor
 from typing import List, Dict
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from torch import FloatTensor
 
 
 class CerebrasMiner(bittensor.BasePromptingMiner):
@@ -83,7 +84,7 @@ def __init__(self):
         )
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass
 

diff --git a/neurons/text/prompting/miners/cohere/neuron.py b/neurons/text/prompting/miners/cohere/neuron.py
@@ -21,6 +21,7 @@
 import bittensor
 from typing import List, Dict
 from langchain.llms import Cohere
+from torch import FloatTensor
 
 
 class CohereMiner(bittensor.BasePromptingMiner):
@@ -108,7 +109,7 @@ def __init__(self):
         )
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass
 

diff --git a/neurons/text/prompting/miners/gooseai/neuron.py b/neurons/text/prompting/miners/gooseai/neuron.py
@@ -18,6 +18,8 @@
 import argparse
 import bittensor
 from typing import List, Dict, Any, Optional
+from torch import FloatTensor
+
 from langchain.llms import GooseAI
 
 
@@ -111,7 +113,7 @@ def __init__(self):
         )
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass
 

diff --git a/neurons/text/prompting/miners/gpt4all/neuron.py b/neurons/text/prompting/miners/gpt4all/neuron.py
@@ -18,6 +18,7 @@
 import argparse
 import bittensor
 from typing import List, Dict
+from torch import FloatTensor
 from langchain.llms import GPT4All
 
 
@@ -164,7 +165,7 @@ def __init__(self):
         )
 
     def backward(
-        self, messages: List[Dict[str, str]], response: str, rewards: torch.FloatTensor
+        self, messages: List[Dict[str, str]], response: str, rewards: FloatTensor
     ) -> str:
         pass