65 lines
1.8 KiB
Python
65 lines
1.8 KiB
Python
from flask import Flask, request, jsonify
|
|
import asyncio
|
|
import axios
|
|
|
|
app = Flask(__name__)
|
|
|
|
# Load Markdown files
|
|
def loadMarkdownFiles(directory):
|
|
files = os.listdir(directory)
|
|
markdownFiles = []
|
|
|
|
for file in files:
|
|
if file.endswith('.md'):
|
|
filePath = os.path.join(directory, file)
|
|
with open(filePath, 'r', encoding='utf-8') as f:
|
|
content = f.read()
|
|
markdownFiles.append({'name': file, 'content': content})
|
|
|
|
return markdownFiles
|
|
|
|
# Initialize RAG model and tokenizer
|
|
async def initRagModel():
|
|
from transformers import AutoTokenizer, RagTokenForGeneration
|
|
tokenizer = await AutoTokenizer.from_pretrained('facebook/rag-token-nq')
|
|
model = await RagTokenForGeneration.from_pretrained('facebook/rag-token-nq')
|
|
|
|
return {'tokenizer': tokenizer, 'model': model}
|
|
|
|
# Retrieve relevant information from Markdown files using Ollama API
|
|
async def retrieveInformation(query):
|
|
try:
|
|
config = configparser.ConfigParser()
|
|
config.read('ollama.ini')
|
|
host = config.get('Ollama', 'host')
|
|
port = config.get('Ollama', 'port')
|
|
|
|
response = await axios.post(f'http://{host}:{port}/chat', {'query': query})
|
|
return response.json()['response']
|
|
except (axios.AxiosError, Exception) as error:
|
|
print(f'Error: {error}')
|
|
raise
|
|
|
|
# Chatbot logic
|
|
async def chatbot():
|
|
directory = './notes' # Directory containing Markdown files
|
|
markdownFiles = loadMarkdownFiles(directory)
|
|
ragModel = await initRagModel()
|
|
|
|
print('Chatbot is ready! Ask your questions.')
|
|
|
|
while True:
|
|
query = input().strip()
|
|
if query.lower() == 'exit':
|
|
break
|
|
|
|
try:
|
|
response = await retrieveInformation(query)
|
|
print(f'Chatbot: {response}')
|
|
except Exception as error:
|
|
print(f'Error: {error}')
|
|
|
|
if __name__ == '__main__':
|
|
asyncio.run(chatbot())
|
|
app.run(debug=True)
|