76 lines
2.1 KiB
Python
76 lines
2.1 KiB
Python
|
from flask import Flask, request, jsonify
|
||
|
|
import asyncio
|
||
|
|
import axios
|
||
|
|
|
||
|
|
app = Flask(__name__)
|
||
|
|
|
||
|
|
# Load Markdown files
|
||
|
|
def loadMarkdownFiles(directory):
|
||
|
|
const files = fs.readdirSync(directory);
|
||
|
|
const markdownFiles = [];
|
||
|
|
|
||
|
|
for (const file of files) {
|
||
|
|
if (path.extname(file).toLowerCase() === '.md') {
|
||
|
|
const filePath = path.join(directory, file);
|
||
|
|
const content = fs.readFileSync(filePath, 'utf-8');
|
||
|
|
markdownFiles.push({ name: file, content });
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
return markdownFiles
|
||
|
|
|
||
|
|
# Initialize RAG model and tokenizer
|
||
|
|
async def initRagModel():
|
||
|
|
const tokenizer = await AutoTokenizer.from_pretrained('facebook/rag-token-nq');
|
||
|
|
const model = await RagTokenForGeneration.from_pretrained('facebook/rag-token-nq');
|
||
|
|
|
||
|
|
return { tokenizer, model }
|
||
|
|
|
||
|
|
# Retrieve relevant information from Markdown files using Ollama API
|
||
|
|
async def retrieveInformation(query):
|
||
|
|
try:
|
||
|
|
response = await axios.post('http://localhost:8080/chat', { query });
|
||
|
|
return response.data.response;
|
||
|
|
except (axios.AxiosError, Exception) as error:
|
||
|
|
print(f'Error: {error.message}')
|
||
|
|
raise
|
||
|
|
|
||
|
|
# Chatbot logic
|
||
|
|
async def chatbot():
|
||
|
|
const directory = './notes'; // Directory containing Markdown files
|
||
|
|
const markdownFiles = loadMarkdownFiles(directory);
|
||
|
|
const ragModel = await initRagModel();
|
||
|
|
|
||
|
|
print('Chatbot is ready! Ask your questions.')
|
||
|
|
|
||
|
|
process.stdin.on('data', async (data) => {
|
||
|
|
const query = data.toString().trim();
|
||
|
|
if (query.toLowerCase() === 'exit') {
|
||
|
|
process.exit(0);
|
||
|
|
}
|
||
|
|
|
||
|
|
try:
|
||
|
|
const response = await retrieveInformation(query);
|
||
|
|
print(f'Chatbot: {response}')
|
||
|
|
except Exception as error:
|
||
|
|
print(f'Error: {error.message}')
|
||
|
|
|
||
|
|
# Flask route to handle chat requests
|
||
|
|
@app.route('/chat', methods=['POST'])
|
||
|
|
async def chat():
|
||
|
|
data = request.json
|
||
|
|
query = data.get('query')
|
||
|
|
|
||
|
|
if not query:
|
||
|
|
return jsonify({'error': 'No query provided'}), 400
|
||
|
|
|
||
|
|
try:
|
||
|
|
response = await retrieveInformation(query)
|
||
|
|
return jsonify({'response': response})
|
||
|
|
except Exception as error:
|
||
|
|
return jsonify({'error': str(error)}), 500
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
asyncio.run(chatbot())
|
||
|
|
app.run(debug=True)
|