Update text_generator.py

text_generator.py

@@ -18,14 +18,14 @@ class TextGenerationTool(Tool):
     
     # 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"
+        "distilgpt2": "distilgpt2",  # Smaller model, may work without auth
+        "gpt2-small": "sshleifer/tiny-gpt2",  # Tiny model for testing
+        "opt-125m": "facebook/opt-125m",  # Small, open model
+        "bloom-560m": "bigscience/bloom-560m",
+        "gpt2": "gpt2"  # Original GPT-2
     }
     
-    def __init__(self, default_model="gpt2", use_api=False):
+    def __init__(self, default_model="distilgpt2", use_api=False):
         """Initialize with a default model and API preference."""
         super().__init__()
         self.default_model = default_model
@@ -33,9 +33,9 @@ class TextGenerationTool(Tool):
         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.")
+        self.token = os.environ.get('HF_TOKEN') or os.environ.get('HF_token')
+        if self.token is None:
+            print("Warning: No Hugging Face token found. Set HF_TOKEN environment variable for authenticated requests.")
     
     def forward(self, text: str):
         """Process the input prompt and generate text."""
@@ -56,31 +56,45 @@ class TextGenerationTool(Tool):
     
     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
+        try:
+            # Get or create the pipeline
+            if model_name not in self._pipelines:
+                # Use token if available, otherwise try without it
+                try:
+                    kwargs = {"token": self.token} if self.token else {}
+                    self._pipelines[model_name] = pipeline(
+                        "text-generation", 
+                        model=model_name, 
+                        **kwargs
+                    )
+                except Exception as e:
+                    print(f"Error loading model {model_name}: {str(e)}")
+                    # Fall back to tiny-distilgpt2 if available
+                    if model_name != "sshleifer/tiny-gpt2":
+                        print("Falling back to tiny-gpt2 model...")
+                        return self._generate_via_pipeline(prompt, "sshleifer/tiny-gpt2", max_length, temperature)
+                    else:
+                        raise e
+            
+            generator = self._pipelines[model_name]
+            
+            # Generate text
+            result = generator(
+                prompt, 
+                max_length=max_length, 
+                num_return_sequences=1, 
+                temperature=temperature
             )
         
-        generator = self._pipelines[model_name]
-        
-        # Generate text
-        result = generator(
-            prompt, 
-            max_length=max_length, 
-            num_return_sequences=1, 
-            temperature=temperature
-        )
-        
-        # Extract and return the 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)
+            # Extract and return the 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)
+        except Exception as e:
+            return f"Error generating text: {str(e)}\n\nPlease try a different model or prompt."
     
     def _generate_via_api(self, prompt, model_name):
         """Generate text by calling the Hugging Face API."""