71 lines
2.0 KiB
JavaScript
71 lines
2.0 KiB
JavaScript
const { AutoTokenizer, RagTokenForGeneration } = require('@huggingface/transformers');
|
|
const fs = require('fs');
|
|
const path = require('path');
|
|
const axios = require('axios');
|
|
const configparser = require('configparser');
|
|
|
|
// Load Markdown files
|
|
function 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 function 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 function retrieveInformation(query) {
|
|
try {
|
|
const config = new configparser.ConfigParser();
|
|
config.read('ollama.ini');
|
|
const host = config.get('Ollama', 'host');
|
|
const port = config.get('Ollama', 'port');
|
|
|
|
const response = await axios.post(`http://${host}:${port}/chat`, { query });
|
|
return response.data.response;
|
|
} catch (error) {
|
|
console.error('Error:', error.message);
|
|
throw error;
|
|
}
|
|
}
|
|
|
|
// Chatbot logic
|
|
async function chatbot() {
|
|
const directory = './notes'; // Directory containing Markdown files
|
|
const markdownFiles = loadMarkdownFiles(directory);
|
|
const ragModel = await initRagModel();
|
|
|
|
console.log('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);
|
|
console.log(`Chatbot: ${response}`);
|
|
} catch (error) {
|
|
console.error('Error:', error.message);
|
|
}
|
|
});
|
|
}
|
|
|
|
chatbot();
|