Chatmate / src /components /textprocess.py
UMESH266
Adding files
af0c0e2
raw
history blame
1.24 kB
# Text generation
from src.logger.logger import logging
from src.exception.exception import customexception
import os
from dotenv import load_dotenv
load_dotenv()
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):
logging.info("Text response generated.")
return self.llm.invoke(input_text)