From a9954b69c05859a0f898415f47213d8bb5cba5c7 Mon Sep 17 00:00:00 2001
From: Cameron Fairchild <cameron@opentensor.dev>
Date: Tue, 25 Jul 2023 21:13:46 -0400
Subject: [PATCH] fix torch typehint on some neurons BT-1329 (#1460)

---
 neurons/text/prompting/miners/AI21/neuron.py       | 3 ++-
 neurons/text/prompting/miners/AlephAlpha/neuron.py | 3 ++-
 neurons/text/prompting/miners/cerebras/neuron.py   | 3 ++-
 neurons/text/prompting/miners/cohere/neuron.py     | 3 ++-
 neurons/text/prompting/miners/gooseai/neuron.py    | 4 +++-
 neurons/text/prompting/miners/gpt4all/neuron.py    | 3 ++-
 6 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/neurons/text/prompting/miners/AI21/neuron.py b/neurons/text/prompting/miners/AI21/neuron.py
index d408468d85..fc346dde28 100644
--- 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
index 0053e67822..ac0f5f11a8 100644
--- 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
index d4665ae985..13e758d327 100644
--- 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
index fb86ad907b..f619c555b4 100644
--- 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
index 967a59c05e..893f0aebb6 100644
--- 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
index be78c08d79..0b98090295 100644
--- 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