Spaces:
Sleeping
Sleeping
File size: 2,628 Bytes
fc6e3b5 3995ebc fc6e3b5 3995ebc fc6e3b5 3995ebc fc6e3b5 3995ebc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
from gtts import gTTS
import base64
import os
from haystack import Document
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.builders import PromptBuilder
from haystack.components.generators.hugging_face_local import HuggingFaceLocalGenerator
from haystack.pipeline import Pipeline
def init_doc_store(path, files):
docs = []
for file in files:
with open(path + '/' + file, 'r') as f:
content = f.read()
docs.append(Document(content=content, meta={'name':file}))
document_store = InMemoryDocumentStore()
document_store.write_documents(docs)
return document_store
def define_components(document_store):
retriever = InMemoryBM25Retriever(document_store, top_k=3)
template = """
Given the following information, answer the question.
Context:
{% for document in documents %}
{{ document.content }}
{% endfor %}
Question: {{question}}
Answer:
"""
prompt_builder = PromptBuilder(template=template)
generator = HuggingFaceLocalGenerator(model="gpt2",
task="text-generation",
# device='cuda',
generation_kwargs={
"max_new_tokens": 100,
"temperature": 0.9,
})
generator.warm_up()
return retreiver, prompt_builder, generator
def define_pipeline(retreiver, prompt_builder, generator):
basic_rag_pipeline = Pipeline()
basic_rag_pipeline.add_component("retriever", retriever)
basic_rag_pipeline.add_component("prompt_builder", prompt_builder)
basic_rag_pipeline.add_component("llm", generator)
basic_rag_pipeline.connect("retriever", "prompt_builder.documents")
basic_rag_pipeline.connect("prompt_builder", "llm")
return basic_rag_pipeline
def generate_response(question, pipeline):
response = pipeline.run({'retriever':{"query":question}, 'prompt_builder':{'question':question}})
response = response['llm']['replies'][0]
return response
def audio_response(response):
audio_stream="response_audio.mp3"
tts = gTTS(response)
tts.save(audio_stream)
with open(audio_stream, 'rb') as file:
audio_data = file.read()
audio_base64 = base64.b64encode(audio_data).decode('utf-8')
audio_tag = f'<audio autoplay="true" src="data:audio/mp3;base64,{audio_base64}">'
return audio_tag
|