Ultima attività 1731289344

knox ha revisionato questo gist 1731289344. Vai alla revisione

1 file changed, 29 insertions

Jina-1.py(file creato)

@@ -0,0 +1,29 @@
1 + from jina import Executor, requests
2 + from docarray import DocList, BaseDoc
3 + from transformers import pipeline
4 +
5 +
6 + class Prompt(BaseDoc):
7 + text: str
8 +
9 +
10 + class Generation(BaseDoc):
11 + prompt: str
12 + text: str
13 +
14 +
15 + class StableLM(Executor):
16 + def __init__(self, **kwargs):
17 + super().__init__(**kwargs)
18 + self.generator = pipeline(
19 + 'text-generation', model='stabilityai/stablelm-base-alpha-3b'
20 + )
21 +
22 + @requests
23 + def generate(self, docs: DocList[Prompt], **kwargs) -> DocList[Generation]:
24 + generations = DocList[Generation]()
25 + prompts = docs.text
26 + llm_outputs = self.generator(prompts)
27 + for prompt, output in zip(prompts, llm_outputs):
28 + generations.append(Generation(prompt=prompt, text=output))
29 + return generations
Più nuovi Più vecchi