myLLM.py

from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain.chains import LanguageModel

class AutoModelLanguageModel(LanguageModel):
    def __init__(self, model_name_or_path):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
        self.model = AutoModelForCausalLM.from_pretrained(model_name_or_path)

    def generate_prompt(self, input_text, history):
        inputs = self.tokenizer.encode(input_text + self.tokenizer.eos_token, return_tensors="pt")
        history = [self.tokenizer.encode(h + self.tokenizer.eos_token, return_tensors="pt") for h in history]
        prompt = torch.cat(history + [inputs], dim=-1)
        return prompt

    def generate_response(self, prompt, max_length):
        output = self.model.generate(prompt, max_length=max_length, pad_token_id=self.tokenizer.pad_token_id)
        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
        return response