Chris4K commited on
Commit
3ff77ad
1 Parent(s): 7de7662

Update text_generator.py

Browse files
Files changed (1) hide show
  1. text_generator.py +2 -1
text_generator.py CHANGED
@@ -16,7 +16,8 @@ class TextGenerationTool(Tool):
16
 
17
  # Initialize the text generation pipeline
18
  #text_generator = pipeline("text-generation") llama mistralai/Mistral-7B-Instruct-v0.1
19
- text_generator = pipeline(model="meta-llama/Llama-2-7b-chat-hf")
 
20
 
21
  # Generate text based on a prompt
22
  generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)
 
16
 
17
  # Initialize the text generation pipeline
18
  #text_generator = pipeline("text-generation") llama mistralai/Mistral-7B-Instruct-v0.1
19
+ text_generator = pipeline(model="gpt2")
20
+ #text_generator = pipeline(model="meta-llama/Llama-2-7b-chat-hf")
21
 
22
  # Generate text based on a prompt
23
  generated_text = text_generator(prompt, max_length=500, num_return_sequences=1, temperature=0.7)