Update utils.py

utils.py

@@ -10,6 +10,7 @@ from langchain.llms import HuggingFacePipeline
 from sentence_transformers import SentenceTransformer, util
 from langchain.chains.question_answering import load_qa_chain
 from transformers import StoppingCriteria, StoppingCriteriaList
+from transformers import T5Tokenizer, T5ForConditionalGeneration, pipeline
 
 
 hf_auth = os.getenv('hf_auth')
@@ -18,25 +19,23 @@ model_id = 'google-t5/t5-base'
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-model_config = transformers.AutoConfig.from_pretrained(
-    model_id,
-    use_auth_token=hf_auth
-)
-
-llm_model = transformers.AutoModelForCausalLM.from_pretrained(
-    model_id,
-    trust_remote_code=True,
-    config=model_config,
-    device_map='auto',
-    use_auth_token=hf_auth
-)
+tokenizer = T5Tokenizer.from_pretrained(model_id)
 
-# enable evaluation mode to allow model inference
-model.eval()
+# Load model
+llm_model = T5ForConditionalGeneration.from_pretrained(model_id)
+llm_model.to(device)
+llm_model.eval()  # Set model to evaluation mode
 
-tokenizer = transformers.AutoTokenizer.from_pretrained(
-    model_id,
-    use_auth_token=hf_auth
+# Define the text generation pipeline
+generate_text = pipeline(
+    'text2text-generation',
+    model=llm_model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1,
+    return_full_text=True,  # Ensure the full text is returned
+    temperature=0.1,        # Control the randomness of the output
+    max_length=512,         # Maximum length of the generated sequence
+    repetition_penalty=1.1  # Penalty to prevent repetition
 )
 
 stop_list = ['\nHuman:', '\n```\n']