gemini gains IQ (embedding fix)

This commit is contained in:
youngdev 2024-11-02 22:24:23 +01:00
parent 32706b00f7
commit 3e6e73f6a1

View file

@ -6,6 +6,28 @@ export class Gemini {
constructor(model_name, url) {
this.model_name = model_name;
this.url = url;
this.safetySettings = [
{
"category": "HARM_CATEGORY_DANGEROUS",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HARASSMENT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_HATE_SPEECH",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
"threshold": "BLOCK_NONE",
},
{
"category": "HARM_CATEGORY_DANGEROUS_CONTENT",
"threshold": "BLOCK_NONE",
},
];
this.genAI = new GoogleGenerativeAI(getKey('GEMINI_API_KEY'));
}
@ -14,12 +36,14 @@ export class Gemini {
let model;
if (this.url) {
model = this.genAI.getGenerativeModel(
{model: this.model_name || "gemini-pro"},
{baseUrl: this.url}
{ model: this.model_name || "gemini-1.5-flash" },
{ baseUrl: this.url },
{ safetySettings: this.safetySettings }
);
} else {
model = this.genAI.getGenerativeModel(
{model: this.model_name || "gemini-pro"}
{ model: this.model_name || "gemini-1.5-flash" },
{ safetySettings: this.safetySettings }
);
}
@ -39,16 +63,16 @@ export class Gemini {
let model;
if (this.url) {
model = this.genAI.getGenerativeModel(
{model: this.model_name || "embedding-001"},
{baseUrl: this.url}
{ model: "text-embedding-004" },
{ baseUrl: this.url }
);
} else {
model = this.genAI.getGenerativeModel(
{model: this.model_name || "embedding-001"}
{ model: "text-embedding-004" }
);
}
const result = await model.embedContent(text);
return result.embedding;
return result.embedding.values;
}
}