feat(backend): Backend config convention and pull model

backend/__tests__/loadAllChatsModels.spec.ts

@@ -0,0 +1,72 @@
+import { ConfigLoader } from '../src/config/ConfigLoader';
+import { ModelDownloader, getModel } from '../src/model/ModelDownloader';
+
+const originalIsArray = Array.isArray;
+
+Array.isArray = jest.fn((type: any): type is any[] => {
+  if (
+    type &&
+    type.constructor &&
+    (type.constructor.name === 'Float32Array' ||
+      type.constructor.name === 'BigInt64Array')
+  ) {
+    return true;
+  }
+  return originalIsArray(type);
+}) as unknown as (arg: any) => arg is any[];
+
+jest.mock('../src/config/ConfigLoader', () => {
+  return {
+    ConfigLoader: jest.fn().mockImplementation(() => {
+      return {
+        get: jest.fn().mockReturnValue({
+          chat1: {
+            model: 'Xenova/LaMini-Flan-T5-783M',
+            task: 'text2text-generation',
+          },
+        }),
+        validateConfig: jest.fn(),
+      };
+    }),
+  };
+});
+
+describe('loadAllChatsModels with real model loading', () => {
+  beforeAll(async () => {
+    await ModelDownloader.downloadAllModels();
+  });
+
+  it('should load real models specified in config', async () => {
+    expect(ConfigLoader).toHaveBeenCalled();
+
+    const chat1Model = getModel('chat1');
+    expect(chat1Model).toBeDefined();
+    console.log('Loaded Model:', chat1Model);
+
+    expect(chat1Model).toHaveProperty('model');
+    expect(chat1Model).toHaveProperty('tokenizer');
+
+    try {
+      const chat1Output = await chat1Model(
+        'Write me a love poem about cheese.',
+        {
+          max_new_tokens: 200,
+          temperature: 0.9,
+          repetition_penalty: 2.0,
+          no_repeat_ngram_size: 3,
+        },
+      );
+
+      console.log('Model Output:', chat1Output);
+      expect(chat1Output).toBeDefined();
+      expect(chat1Output[0]).toHaveProperty('generated_text');
+      console.log(chat1Output[0].generated_text);
+    } catch (error) {
+      console.error('Error during model inference:', error);
+    }
+  }, 60000);
+});
+
+afterAll(() => {
+  Array.isArray = originalIsArray;
+});

backend/package.json

@@ -26,6 +26,8 @@
   },
   "dependencies": {
     "@apollo/server": "^4.11.0",
+    "@huggingface/hub": "latest",
+    "@huggingface/transformers": "latest",
     "@nestjs/apollo": "^12.2.0",
     "@nestjs/axios": "^3.0.3",
     "@nestjs/common": "^10.0.0",
@@ -36,11 +38,13 @@
     "@nestjs/platform-express": "^10.0.0",
     "@nestjs/typeorm": "^10.0.2",
     "@types/bcrypt": "^5.0.2",
+    "axios": "^1.7.7",
     "@types/fs-extra": "^11.0.4",
     "bcrypt": "^5.1.1",
     "class-validator": "^0.14.1",
     "fs-extra": "^11.2.0",
     "graphql": "^16.9.0",
+    "lodash": "^4.17.21",
     "graphql-subscriptions": "^2.0.0",
     "graphql-ws": "^5.16.0",
     "markdown-to-txt": "^2.0.1",

backend/src/config/ConfigLoader.ts

@@ -0,0 +1,48 @@
+import * as fs from 'fs';
+import * as path from 'path';
+import _ from 'lodash';
+export interface ChatConfig {
+  model: string;
+  endpoint?: string;
+  token?: string;
+  default?: boolean;
+  task?: string;
+}
+export class ConfigLoader {
+  private config: ChatConfig;
+
+  private readonly configPath: string;
+
+  constructor() {
+    this.configPath = path.resolve(__dirname, 'config.json');
+    this.loadConfig();
+  }
+
+  private loadConfig() {
+    const file = fs.readFileSync(this.configPath, 'utf-8');
+    this.config = JSON.parse(file);
+  }
+
+  get<T>(path: string) {
+    return _.get(this.config, path);
+  }
+
+  set(path: string, value: any) {
+    _.set(this.config, path, value);
+    this.saveConfig();
+  }
+
+  private saveConfig() {
+    fs.writeFileSync(
+      this.configPath,
+      JSON.stringify(this.config, null, 4),
+      'utf-8',
+    );
+  }
+
+  validateConfig() {
+    if (!this.config) {
+      throw new Error("Invalid configuration: 'chats' section is missing.");
+    }
+  }
+}

backend/src/model/ModelDownloader.ts

@@ -0,0 +1,63 @@
+import { Logger } from '@nestjs/common';
+import { PipelineType, pipeline, env } from '@huggingface/transformers';
+import { ConfigLoader, ChatConfig } from '../config/ConfigLoader';
+
+type Progress = { loaded: number; total: number };
+type ProgressCallback = (progress: Progress) => void;
+
+env.allowLocalModels = false;
+
+export class ModelDownloader {
+  private static readonly logger = new Logger(ModelDownloader.name);
+  private static readonly loadedModels = new Map<string, any>();
+
+  public static async downloadAllModels(
+    progressCallback: ProgressCallback = () => {},
+  ): Promise<void> {
+    const configLoader = new ConfigLoader();
+    configLoader.validateConfig();
+    const chats = configLoader.get<{ [key: string]: ChatConfig }>('chats');
+
+    const loadPromises = Object.entries(chats).map(
+      async ([chatKey, chatConfig]: [string, ChatConfig]) => {
+        const { model, task } = chatConfig;
+        try {
+          ModelDownloader.logger.log(`Starting to load model: ${model}`);
+          const pipelineInstance = await ModelDownloader.downloadModel(
+            task,
+            model,
+            progressCallback,
+          );
+          ModelDownloader.logger.log(`Model loaded successfully: ${model}`);
+          this.loadedModels.set(chatKey, pipelineInstance);
+        } catch (error) {
+          ModelDownloader.logger.error(
+            `Failed to load model ${model}:`,
+            error.message,
+          );
+        }
+      },
+    );
+
+    await Promise.all(loadPromises);
+
+    ModelDownloader.logger.log('All models loaded.');
+  }
+
+  private static async downloadModel(
+    task: string,
+    model: string,
+    progressCallback?: ProgressCallback,
+  ): Promise<any> {
+    const pipelineOptions = progressCallback
+      ? { progress_callback: progressCallback }
+      : undefined;
+    return pipeline(task as PipelineType, model, pipelineOptions);
+  }
+
+  public static getModel(chatKey: string): any {
+    return ModelDownloader.loadedModels.get(chatKey);
+  }
+}
+
+export const getModel = ModelDownloader.getModel;