| from langchain.llms import HuggingFacePipeline | |
| from langchain.chains import LLMChain | |
| from langchain.prompts import PromptTemplate | |
| from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer | |
| class NameExtractorChain: | |
| def __init__(self, model_name: str = "name-extraction"): | |
| self.pipe = pipeline( | |
| "text2text-generation", | |
| model=model_name, | |
| max_new_tokens=10, | |
| model_kwargs={"temperature": 0} | |
| ) | |
| self.llm = HuggingFacePipeline(pipeline=self.pipe) | |
| self.prompt = PromptTemplate( | |
| input_variables=["conversation"], | |
| template="""Extract only the name of the person from this conversation. | |
| If there's no name, return 'No name found'. | |
| Conversation: {conversation}""") | |
| self.chain = LLMChain(llm=self.llm, prompt=self.prompt) | |
| def extract_name(self, text: str): | |
| text=text.strip() | |
| if len(text.split())==1: | |
| text= "It's " + text | |
| try: | |
| output = self.chain.run(conversation=text) | |
| return output | |
| except Exception as e: | |
| print(f"Error processing text: {str(e)}") | |
| extractor = NameExtractorChain() | |
| print(extractor.extract_name(sample-text)) | |