CTX-5783: Additional linter errors fixed

coretex-ai · Jul 4, 2024 · 3a272f5 · 3a272f5
1 parent 8e2906d
commit 3a272f5
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 15 deletions.
diff --git a/tasks/stable-diffusion/main.py b/tasks/stable-diffusion/main.py
@@ -18,9 +18,9 @@
 
 def loadModel(device: str) -> StableDiffusionPipeline:
     dtype = torch.float16 if device == "cuda" else torch.float32
-    pipe: StableDiffusionPipeline = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype = dtype)
-    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
-    return pipe.to(device)
+    pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype = dtype)  # type: ignore[no-untyped-call]
+    pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)  # type: ignore[no-untyped-call]
+    return pipe.to(device)  # type: ignore[no-any-return]
 
 
 def getDefault(taskRun: TaskRun, name: str, default: Any) -> Any:
@@ -53,7 +53,7 @@ def generateImages(
         # Create an array equal to number of input prompts
         negativePrompts = [negativePrompt] * len(prompts)
 
-    images = model(
+    images = model(  # type: ignore[operator]
         prompts,
         negative_prompt = negativePrompts,
         num_inference_steps = steps,

diff --git a/tasks/translation-ollama/main.py b/tasks/translation-ollama/main.py
@@ -13,7 +13,7 @@
 
 def readPDF(filePath: Path) -> list[str]:
     pagesText: list[str] = []
-    
+
     with fitz.open(filePath) as doc:
         for page in doc:
             paragraphs = page.get_text().split("\n")
@@ -24,14 +24,14 @@ def readPDF(filePath: Path) -> list[str]:
 
 def loadCorpus(dataset: CustomDataset) -> list[list[str]]:
     corpus: list[list[str]] = []
-    
+
     for sample in dataset.samples:
         sample.unzip()
 
         pdfPaths = list(sample.path.rglob("*.pdf"))
         if len(pdfPaths) == 0:
             raise ValueError(">> [LLM Translate] The provided dataset does not contain any .pdf documents")
-        
+
         for pdfPath in pdfPaths:
             if not "__MACOSX" in str(pdfPath):
                 corpus.append(readPDF(pdfPath))
@@ -45,21 +45,21 @@ def main() -> None:
     dataset.download()
 
     launchOllamaServer()
-    
+
     logging.info(">> [OllamaRAG] Pulling model")
     pullModel(LLM)
-    
+
     logging.info(">> [OllamaRAG] Loading text corpus")
     corpus = loadCorpus(taskRun.dataset)
-    
+
     translatedDataset = CustomDataset.createDataset(f"{taskRun.id}-translated", taskRun.projectId)
 
     language = taskRun.parameters["language"]
 
     for counter, document in enumerate(corpus, start = 1):
         document = [x.strip() for x in document]
         document = [line for line in document if line != ""]
-        
+
         translatedText = ""
         for paragraph in document:
             logging.info(">> [OllamaRAG] Translating paragraph")
@@ -70,10 +70,10 @@ def main() -> None:
                 "role": "user",
                 "content": query
             }
-            response = ollama.chat(model = LLM, messages = [msg])
-            answer = response["message"]["content"]
+            response = ollama.chat(model = LLM, messages = [msg])  # type: ignore[list-item]
+            answer = response["message"]["content"]  # type: ignore[index]
             translatedText += answer + "\n"
-        
+
         txtFileName = f"file-{counter}.txt"
         txtFile = folder_manager.temp / txtFileName
         with open(txtFile, "w") as f:
@@ -83,7 +83,7 @@ def main() -> None:
         zipFile = folder_manager.temp / zipFileName
         with zipfile.ZipFile(zipFile, "w") as zf:
             zf.write(txtFile, txtFileName)
-        
+
         translatedDataset.add(zipFile)
 
         taskRun.submitOutput("translatedDataset", translatedDataset)