Update text_generator.py

text_generator.py

@@ -1,68 +1,116 @@
-import requests
 import os
-from transformers import pipeline
-
-
-from transformers import Tool
-# Import other necessary libraries if needed
+import requests
+import gradio as gr
+from transformers import pipeline, Tool
 
 class TextGenerationTool(Tool):
     name = "text_generator"
-    description = (
-        "This is a tool for text generation. It takes a prompt as input and returns the generated text."
-    )
-
+    description = "This is a tool for text generation. It takes a prompt as input and returns the generated text."
+    
     inputs = ["text"]
     outputs = ["text"]
-
+    
+    # Available text generation models
+    models = {
+        "orca": "microsoft/Orca-2-13b",
+        "gpt2-dolly": "lgaalves/gpt2-dolly",
+        "gpt2": "gpt2",
+        "bloom": "bigscience/bloom-560m",
+        "openchat": "openchat/openchat_3.5"
+    }
+    
+    def __init__(self, default_model="orca", use_api=False):
+        """Initialize with a default model and API preference."""
+        self.default_model = default_model
+        self.use_api = use_api
+        self._pipelines = {}
+        
+        # Check for API token
+        self.token = os.environ.get('HF_token')
+        if self.token is None and use_api:
+            print("Warning: HF_token environment variable not set. API calls will fail.")
+    
     def __call__(self, prompt: str):
-        #API_URL = "https://api-inference.huggingface.co/models/openchat/openchat_3.5"
-        #headers = {"Authorization": "Bearer " +  os.environ['hf']}
-        token=os.environ['HF_token']
-        #payload = {
-        #    "inputs": prompt  # Adjust this based on your model's input format
-        #}
-
-        #payload = {
-        #            "inputs": "Can you please let us know more details about your ",
-        #        }
+        """Process the input prompt and generate text."""
+        return self.generate_text(prompt)
+    
+    def generate_text(self, prompt, model_key=None, max_length=500, temperature=0.7):
+        """Generate text based on the prompt using the specified or default model."""
+        # Determine which model to use
+        model_key = model_key or self.default_model
+        model_name = self.models.get(model_key, self.models[self.default_model])
         
-        #def query(payload):
-        #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="lgaalves/gpt2-dolly", token=token)
-        text_generator = pipeline(model="microsoft/Orca-2-13b", token=token)
-
-        # 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
+        # Generate using API if specified
+        if self.use_api and model_key == "openchat":
+            return self._generate_via_api(prompt, model_name)
+        
+        # Otherwise use local pipeline
+        return self._generate_via_pipeline(prompt, model_name, max_length, temperature)
+    
+    def _generate_via_pipeline(self, prompt, model_name, max_length, temperature):
+        """Generate text using a local pipeline."""
+        # Get or create the pipeline
+        if model_name not in self._pipelines:
+            self._pipelines[model_name] = pipeline(
+                "text-generation", 
+                model=model_name, 
+                token=self.token
+            )
+        
+        generator = self._pipelines[model_name]
+        
+        # Generate text
+        result = generator(
+            prompt, 
+            max_length=max_length, 
+            num_return_sequences=1, 
+            temperature=temperature
+        )
         
-        # Define the payload for the request
-        #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()
-
         # Extract and return the generated text
-        #return generated_text["generated_text"]
-
-# Uncomment and customize the following lines based on your text generation needs
-# text_generator = pipeline(model="gpt2")
-# generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
-
-# Print the generated text if needed
-# print(generated_text)
+        if isinstance(result, list) and len(result) > 0:
+            if isinstance(result[0], dict) and 'generated_text' in result[0]:
+                return result[0]['generated_text']
+            return result[0]
+        
+        return str(result)
+    
+    def _generate_via_api(self, prompt, model_name):
+        """Generate text by calling the Hugging Face API."""
+        if not self.token:
+            return "Error: HF_token not set. Cannot use API."
+        
+        api_url = f"https://api-inference.huggingface.co/models/{model_name}"
+        headers = {"Authorization": f"Bearer {self.token}"}
+        payload = {"inputs": prompt}
+        
+        try:
+            response = requests.post(api_url, headers=headers, json=payload)
+            response.raise_for_status()  # Raise exception for HTTP errors
+            
+            result = response.json()
+            
+            # Handle different response formats
+            if isinstance(result, list) and len(result) > 0:
+                if isinstance(result[0], dict) and 'generated_text' in result[0]:
+                    return result[0]['generated_text']
+            elif isinstance(result, dict) and 'generated_text' in result:
+                return result['generated_text']
+            
+            # Fall back to returning the raw response
+            return str(result)
+            
+        except Exception as e:
+            return f"Error generating text: {str(e)}"
+
+# For standalone testing
+if __name__ == "__main__":
+    # Create an instance of the TextGenerationTool
+    text_generator = TextGenerationTool(default_model="gpt2")
+    
+    # Test with a simple prompt
+    test_prompt = "Once upon a time in a digital world,"
+    result = text_generator(test_prompt)
+    
+    print(f"Prompt: {test_prompt}")
+    print(f"Generated text:\n{result}")