jbilcke-hf HF staff commited on
Commit
2f7798c
β€’
1 Parent(s): a5a608f

try to improve robustness of LLM responses

Browse files
src/app/queries/getStory.ts CHANGED
@@ -1,10 +1,11 @@
1
  import { createLlamaPrompt } from "@/lib/createLlamaPrompt"
 
 
 
2
 
3
  import { predict } from "./predict"
4
  import { Preset } from "../engine/presets"
5
 
6
- type LLMResponse = Array<{panel: number; caption: string }>
7
-
8
  export const getStory = async ({
9
  preset,
10
  prompt = "",
@@ -32,6 +33,7 @@ export const getStory = async ({
32
 
33
 
34
  let result = ""
 
35
  try {
36
  result = await predict(query)
37
  if (!result.trim().length) {
@@ -51,33 +53,45 @@ export const getStory = async ({
51
  }
52
 
53
  console.log("Raw response from LLM:", result)
54
- let tmp = result // result.split("Caption:").pop() || result
55
- tmp = tmp
56
- .replaceAll("}}", "}")
57
- .replaceAll("]]", "]")
58
- .replaceAll(",,", ",")
59
 
60
  try {
61
- // we only keep what's after the first [
62
- let jsonOrNot = `[${tmp.split("[").pop() || ""}`
63
-
64
- // and before the first ]
65
- jsonOrNot = `${jsonOrNot.split("]").shift() || ""}]`
66
-
67
- const jsonData = JSON.parse(jsonOrNot) as LLMResponse
68
- const captions = jsonData.map(item => item.caption.trim())
69
- return captions.map(caption => caption.split(":").pop()?.trim() || "")
70
  } catch (err) {
71
  console.log(`failed to read LLM response: ${err}`)
72
 
73
- // in case of failure, it might be because the LLM hallucinated a completely different response,
74
- // such as markdown. There is no real solution.. but we can try a fallback:
 
 
 
 
 
75
 
76
- const candidateList = (
77
- tmp.split("*")
78
- .map(item => item.replaceAll("[", "[").replaceAll("]", "]").trim())
79
- )
 
 
 
 
 
80
 
81
- return candidateList
 
 
 
 
 
 
 
 
 
 
82
  }
 
 
83
  }
 
1
  import { createLlamaPrompt } from "@/lib/createLlamaPrompt"
2
+ import { dirtyLLMResponseCleaner } from "@/lib/dirtyLLMResponseCleaner"
3
+ import { dirtyLLMJsonParser } from "@/lib/dirtyLLMJsonParser"
4
+ import { dirtyCaptionCleaner } from "@/lib/dirtyCaptionCleaner"
5
 
6
  import { predict } from "./predict"
7
  import { Preset } from "../engine/presets"
8
 
 
 
9
  export const getStory = async ({
10
  preset,
11
  prompt = "",
 
33
 
34
 
35
  let result = ""
36
+
37
  try {
38
  result = await predict(query)
39
  if (!result.trim().length) {
 
53
  }
54
 
55
  console.log("Raw response from LLM:", result)
56
+ const tmp = dirtyLLMResponseCleaner(result)
57
+
58
+ let captions: string[] = []
 
 
59
 
60
  try {
61
+ captions = dirtyLLMJsonParser(tmp)
 
 
 
 
 
 
 
 
62
  } catch (err) {
63
  console.log(`failed to read LLM response: ${err}`)
64
 
65
+ // it is possible that the LLM has generated multiple JSON files like this:
66
+
67
+ /*
68
+ [ {
69
+ "panel": 1,
70
+ "caption": "A samurai stands at the edge of a bustling street in San Francisco, looking out of place among the hippies and beatniks."
71
+ } ]
72
 
73
+ [ {
74
+ "panel": 2,
75
+ "caption": "The samurai spots a group of young people playing music on the sidewalk. He approaches them, intrigued."
76
+ } ]
77
+ */
78
+ try {
79
+ // in that case, we can try to repair it like so:
80
+ let strategy2 = `[${tmp.split("[").pop() || ""}`
81
+ strategy2.replaceAll("[", ",")
82
 
83
+ captions = dirtyLLMJsonParser(strategy2)
84
+ } catch (err2) {
85
+
86
+ // in case of failure here, it might be because the LLM hallucinated a completely different response,
87
+ // such as markdown. There is no real solution.. but we can try a fallback:
88
+
89
+ captions = (
90
+ tmp.split("*")
91
+ .map(item => item.replaceAll("[", "[").replaceAll("]", "]").trim())
92
+ )
93
+ }
94
  }
95
+
96
+ return captions.map(caption => dirtyCaptionCleaner(caption))
97
  }
src/lib/dirtyCaptionCleaner.ts ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ export function dirtyCaptionCleaner(input: string) {
2
+ return input.split(":").pop()?.trim() || ""
3
+ }
src/lib/dirtyLLMJsonParser.ts ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import { LLMResponse } from "@/types"
2
+
3
+ export function dirtyLLMJsonParser(input: string): string[] {
4
+ // we only keep what's after the first [
5
+ let jsonOrNot = `[${input.split("[").pop() || ""}`
6
+
7
+ // and before the first ]
8
+ jsonOrNot = `${jsonOrNot.split("]").shift() || ""}]`
9
+
10
+ const jsonData = JSON.parse(jsonOrNot) as LLMResponse
11
+
12
+ const captions = jsonData.map(item => item.caption.trim())
13
+
14
+ return captions
15
+ }
src/lib/dirtyLLMResponseCleaner.ts ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ export function dirtyLLMResponseCleaner(input: string) {
2
+ return (
3
+ `${input || ""}`
4
+ .replaceAll("}}", "}")
5
+ .replaceAll("]]", "]")
6
+ .replaceAll(",,", ",")
7
+ )
8
+ }
src/types.ts CHANGED
@@ -78,3 +78,5 @@ export interface ImageAnalysisResponse {
78
  result: string
79
  error?: string
80
  }
 
 
 
78
  result: string
79
  error?: string
80
  }
81
+
82
+ export type LLMResponse = Array<{panel: number; caption: string }>