Update text_generator.py

text_generator.py

@@ -14,23 +14,37 @@ class TextGenerationTool(Tool):
     outputs = ["text"]
 
     def __call__(self, prompt: str):
-        API_URL = "https://api-inference.huggingface.co/models/openchat/openchat_3.5"
-        headers = {"Authorization": "Bearer " +  os.environ['hf']}
+        #API_URL = "https://api-inference.huggingface.co/models/openchat/openchat_3.5"
+        #headers = {"Authorization": "Bearer " +  os.environ['hf']}
 
-        payload = {
-            "inputs": prompt  # Adjust this based on your model's input format
-        }
+        #payload = {
+        #    "inputs": prompt  # Adjust this based on your model's input format
+        #}
 
-        payload = {
-                    "inputs": "Can you please let us know more details about your ",
-                }
+        #payload = {
+        #            "inputs": "Can you please let us know more details about your ",
+        #        }
         
         #def query(payload):
-        generated_text = requests.post(API_URL, headers=headers, json=payload).json()
-        print(generated_text)
-        return generated_text["text"]
+        #generated_text = requests.post(API_URL, headers=headers, json=payload).json()
+        #print(generated_text)
+        #return generated_text["text"]
             
+        # Replace the following line with your text generation logic
+        #generated_text = f"Generated text based on the prompt: '{prompt}'"
+
+        # Initialize the text generation pipeline
+        text_generator = pipeline(model="bigcode/starcoder")
+
+        # Generate text based on a prompt
+        generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
+
+        # Print the generated text
+        print(generated_text)
+
+
 
+        return generated_text
         
         # Define the payload for the request
         #payload = {