File size: 1,656 Bytes
852dc0d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import { createContext, useContext, useState, ReactNode } from 'react'
import { ChatMessage } from '../types'
export interface GenerationConfigState {
temperature: number
maxTokens: number
topP: number
topK: number
doSample: boolean
}
interface TextGenerationContextType {
config: GenerationConfigState
setConfig: React.Dispatch<React.SetStateAction<GenerationConfigState>>
messages: ChatMessage[]
setMessages: React.Dispatch<React.SetStateAction<ChatMessage[]>>
updateSystemMessage: (content: string) => void
}
const TextGenerationContext = createContext<
TextGenerationContextType | undefined
>(undefined)
export function TextGenerationProvider({ children }: { children: ReactNode }) {
const [config, setConfig] = useState<GenerationConfigState>({
temperature: 0.7,
maxTokens: 100,
topP: 0.9,
topK: 50,
doSample: true
})
const [messages, setMessages] = useState<ChatMessage[]>([
{ role: 'system', content: 'You are a helpful assistant.' }
])
const updateSystemMessage = (content: string) => {
setMessages((prev) => [
{ role: 'system', content },
...prev.filter((msg) => msg.role !== 'system')
])
}
const value = {
config,
setConfig,
messages,
setMessages,
updateSystemMessage
}
return (
<TextGenerationContext.Provider value={value}>
{children}
</TextGenerationContext.Provider>
)
}
export function useTextGeneration() {
const context = useContext(TextGenerationContext)
if (context === undefined) {
throw new Error(
'useTextGeneration must be used within a TextGenerationProvider'
)
}
return context
}
|