Update tokenizer.py

tokenizer.py

@@ -1,56 +1,41 @@
+import os
 from transformers import pipeline
-from huggingface_hub import hf_hub_url, cached_download, HfApi
-
-HF_TOKEN = os.environ.get("HF_TOKEN", None)
-
-def generate_code_from_model(input_text):
-    """Generates code using the bigscience/T0_3B model."""
-    model_name = 'bigscience/T0_3B'  # Choose your model
-    generator = pipeline('text-generation', model=model_name)
-    model_type = {
-        "model_type": "compositional",
-        "task": "Develop an app that allows users to search for and modify files on a remote server using the SSH protocol",
-        "history": "",
-        "history_additions": [
-            "Set the task description to 'Develop an app that allows users to search for and modify files on a remote server using the SSH protocol'",
-            "Updated the task description to 'Connect to the remote server using the SSH protocol'",
-            "Completed the task: Connect to the remote server using the SSH protocol",
-            "Updated the task description to 'Search for a file on the remote server'",
-            "Completed the task: Search for a file on the remote server",
-            "Updated the task description to 'Modify the file found on the remote server'",
-            "Completed the task: Modify the file found on the remote server",
-            "Updated the task description to 'Verify the file modifications using the updated file'",
-            "Completed the task: Verify the file modifications using the updated file",
-            "Updated the task description to 'Create a summary of the timeline of progress and any important challenges faced during the development of the app'",
-            "Completed the task: Create a summary of the timeline of progress and any important challenges faced during the development of the app",
-            "Updated the task description to 'Deploy the app to the production server'",
-            "Completed the task: Deploy the app to the production server",
-            "Updated the task description to 'Monitor the performance of the deployed app and make adjustments as necessary'",
-            "Completed the task: Monitor the performance of the deployed app and make adjustments as necessary",
-            "Updated the task description to 'Maintain and update the app in response to user feedback and any discovered bugs'"
-        ],
-        "history_removals": [],
-        "model_log": [
-            "Model prompt: You are attempting to complete the task task: Develop an app that allows users to search for and modify files on a remote server using the SSH protocol Progress:  Connect to the remote server using the SSH protocol",
-            "Model response: Connect to the remote server using the SSH protocol",
-            "Model prompt: Updated the task description to 'Search for a file on the remote server'",
-            "Model response: Search for a file on the remote server",
-            "Model prompt: Completed the task: Search for a file on the remote server Progress:  Connect to the remote server using the SSH protocol\nSearch for a file on the remote server",
-            "Model response: You have successfully completed the task: Search for a file on the remote server",
-            "Model prompt: Updated the task description to 'Modify the file found on the remote server'",
-            "Model response: Modify the file found on the remote server",
-            "Model prompt: Completed the task: Modify the file found on the remote server Progress:  Connect to the remote server using the SSH protocol\nSearch for a file on the remote server\nModify the file found on the remote server",
-            "Model response: You have successfully completed the task: Modify the file found on the remote server",
-            "Model prompt: Updated the task description to 'Verify the file modifications using the updated file'",
-            "Model response: Verify the file modifications using the updated file",
-            "Model prompt: Completed the task: Verify the file modifications using the updated file Progress:  Connect to the remote server using the SSH protocol\nSearch for a file on the remote server\nModify the file found on the remote server\nVerify the file modifications using the updated file",
-            "Model response: You have successfully completed the task: Verify the file modifications using the updated file",
-            "Model prompt: Updated the task description to 'Create a summary of the timeline of progress and any important challenges faced during the development of the app'",
-            "Model response: Create a summary of the timeline of progress and any important challenges faced during the development of the app",
-            "Model prompt: Completed the task: Create a summary of the timeline of progress and any important challenges faced.",
-            "Model response: All tasks complete. What shall I do next?"
-        ]
-    }
-
-    code = generator(input_text, max_length=50, num_return_sequences=1, do_sample=True)[0]['generated_text']
-    return code
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+class CodeGenerator:
+    def __init__(self, model_name='bigscience/T0_3B'):
+        self.generator = pipeline('text-generation', model=model_name)
+
+    def generate_code(self, task_description):
+        """
+        Generates code based on the provided task description using the specified language model.
+
+        Parameters:
+        task_description (str): The task description or prompt for generating the code.
+
+        Returns:
+        str: The generated code.
+        """
+        return self._generate_code_from_model(task_description)
+
+    def _generate_code_from_model(self, input_text):
+        """
+        Internal method to generate code from the model.
+
+        Parameters:
+        input_text (str): The input text for code generation.
+
+        Returns:
+        str: The code generated by the language model.
+        """
+        return self.generator(input_text, max_length=50, num_return_sequences=1, do_sample=True)[0]['generated_text']
+
+def main():
+    task_description = "Develop an app that allows users to search for and modify files on a remote server using the SSH protocol"
+    code_generator = CodeGenerator()
+    generated_code = code_generator.generate_code(task_description)
+    print(generated_code)
+
+if __name__ == "__main__":
+    main()