File size: 1,296 Bytes
af0c0e2 d359a51 af0c0e2 d6f7b0e d359a51 af0c0e2 d359a51 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
# Text generation
from src.logger.logger import logging
from src.exception.exception import customexception
import sys
from langchain_huggingface import HuggingFaceEndpoint
# Text generation model
repo_id="mistralai/Mistral-7B-Instruct-v0.3"
# repo_id='openai-community/gpt2'
class TextProcessor:
def __init__(self, hf_token):
self.llm = HuggingFaceEndpoint(
repo_id=repo_id,
max_new_tokens=512,
top_k=10,
top_p=0.95,
typical_p=0.95,
temperature=0.01,
repetition_penalty=1.03,
streaming=False,
# huggingfacehub_api_token= hf_token,
# stop_sequences=['?', '</s>', '.\n\n']
)
logging.info("LLM model for text generation created.")
def generate_response(self, input_text):
try:
logging.info("Text response generated.")
return self.llm.invoke(input_text)
except Exception as e:
raise customexception(e,sys) |