datasetsANDmodels commited on
Commit
638f5cd
·
verified ·
1 Parent(s): 7b4f98d

Upload usage.py

Browse files
Files changed (1) hide show
  1. usage.py +38 -0
usage.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain.llms import HuggingFacePipeline
2
+ from langchain.chains import LLMChain
3
+ from langchain.prompts import PromptTemplate
4
+ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
5
+
6
+ class NameExtractorChain:
7
+ def __init__(self, model_name: str = "name-extraction"):
8
+ self.pipe = pipeline(
9
+ "text2text-generation",
10
+ model=model_name,
11
+ max_new_tokens=10,
12
+ model_kwargs={"temperature": 0}
13
+ )
14
+
15
+ self.llm = HuggingFacePipeline(pipeline=self.pipe)
16
+
17
+ self.prompt = PromptTemplate(
18
+ input_variables=["conversation"],
19
+
20
+ template="""Extract only the name of the person from this conversation.
21
+ If there's no name, return 'No name found'.
22
+ Conversation: {conversation}""")
23
+ self.chain = LLMChain(llm=self.llm, prompt=self.prompt)
24
+
25
+
26
+ def extract_name(self, text: str):
27
+ text=text.strip()
28
+ if len(text.split())==1:
29
+ text= "It's " + text
30
+ try:
31
+ output = self.chain.run(conversation=text)
32
+ return output
33
+ except Exception as e:
34
+ print(f"Error processing text: {str(e)}")
35
+
36
+
37
+ extractor = NameExtractorChain()
38
+ print(extractor.extract_name(sample-text))