| 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)) |
|
|
