Update crypto_analysis_agents.py

crypto_analysis_agents.py

@@ -1,10 +1,9 @@
 from crewai import Agent
 from langchain.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from crypto_tools import CryptoTools
 from news_tools import NewsTools
 from sentiment_tools import SentimentTools
-import torch
 import os
 
 class CryptoAnalysisAgents:
@@ -14,29 +13,22 @@ class CryptoAnalysisAgents:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(model_name)
 
-        # Function to handle generation with manual configuration
-        def generate_text(input_text, max_new_tokens=512):
-            input_ids = tokenizer.encode(input_text, return_tensors="pt")
-            max_length = input_ids.shape[1] + max_new_tokens  # Ensure max_length accommodates input + generation
-
-            # Generate the text using the model's generate method
-            output_ids = model.generate(
-                input_ids=input_ids,
-                max_length=max_length,  # Set max_length to handle input + generated tokens
-                temperature=0.7,
-                top_p=0.95,
-                top_k=50,
-                repetition_penalty=1.1,
-                do_sample=True,
-                pad_token_id=tokenizer.eos_token_id  # Ensure the model doesn't stop too early
-            )
-
-            # Decode the output ids to get the generated text
-            output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-            return output_text
+        # Create the text-generation pipeline
+        pipe = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=512,  # Control the maximum number of tokens to generate
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+            top_k=50,
+            repetition_penalty=1.1,
+            pad_token_id=tokenizer.eos_token_id  # Ensure correct padding
+        )
 
-        # Wrap this in a LangChain pipeline
-        self.llm = HuggingFacePipeline(generate=generate_text)
+        # Wrap the pipeline in a HuggingFacePipeline
+        self.llm = HuggingFacePipeline(pipeline=pipe)
 
     def market_analyst(self):
         return Agent(
@@ -90,6 +82,6 @@ if __name__ == "__main__":
     # Input text for generation (example)
     input_text = "Analyze the current state of Bitcoin in the market."
     
-    # Generate a report
-    report = market_analyst.llm.generate(input_text)
+    # Generate a report using the pipeline
+    report = market_analyst.llm(input_text)
     print(report)