Update text_generator.py

text_generator.py

@@ -1,5 +1,5 @@
-import os
 import requests
+import os
 
 from transformers import Tool
 # Import other necessary libraries if needed
@@ -14,19 +14,28 @@ class TextGenerationTool(Tool):
     outputs = ["text"]
 
     def __call__(self, prompt: str):
-        API_URL = "https://api-inference.huggingface.co/models/gpt2"
-        headers = {"Authorization": "Bearer " + os.environ['hf']}
-
+        API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
+        headers = {"Authorization": + os.environ['hf']}
+        
+        def query(payload):
+            response = requests.post(API_URL, headers=headers, json=payload)
+            print(response)
+            return response.json()
+            
+        output = query({
+            "inputs": prompt,
+        })
+        
         # Define the payload for the request
-        payload = {
-            "inputs": prompt  # Adjust this based on your model's input format
-        }
+        #payload = {
+        #    "inputs": prompt  # Adjust this based on your model's input format
+        #}
 
         # Make the request to the API
-        generated_text = requests.post(API_URL, headers=headers, json=payload).json()
+        #generated_text = requests.post(API_URL, headers=headers, json=payload).json()
 
         # Extract and return the generated text
-        return generated_text["generated_text"]
+        #return generated_text["generated_text"]
 
 # Uncomment and customize the following lines based on your text generation needs
 # text_generator = pipeline(model="gpt2")