File size: 1,866 Bytes
dbe5e76
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
import io
import json
import random
import requests
import gradio as gr

class TextGenerator:
    def __init__(self, token, model_id):
        self.token = token
        self.model_id = model_id

    def generate_text(self, prompt, hf_model=None, hf_token=None, parameters=None):
        if hf_token is None:
            hf_token = self.token
        if hf_model is None:
            hf_model = self.model_id
        if parameters is None:
            parameters = {'max_new_tokens': 512, 'do_sample': True, 'return_full_text': False, 'temperature': 1.0, 'top_k': 50, 'repetition_penalty': 1.2}
        url = f'https://api-inference.huggingface.co/models/{hf_model}'
        headers = {'Authorization': f'Bearer {hf_token}', 'Content-type': 'application/json'}
        data = {'inputs': prompt, 'stream': False, 'options': {'use_cache': False}, 'parameters': parameters}
        r = requests.post(url, headers=headers, data=json.dumps(data))
        return json.loads(r.content.decode("utf-8"))[0]['generated_text']

class GradioUI:
    def __init__(self, text_generator):
        self.text_generator = text_generator
        self.styles = """..."""  # Add CSS styles here
        self.title_placeholders = ['Pokemon training story']

    def random_title(self):
        return random.choice(self.title_placeholders)

    def generate_interface(self):
        with gr.Blocks(css=self.styles) as demo:
            # Gradio UI setup goes here
            random_title_btn.click(fn=None, inputs=None, outputs=[title], _js=f"return ['{self.random_title()}'];")
            # Additional Gradio configurations go here...
        demo.queue(concurrency_count=5, max_size=256)
        demo.launch()

token = os.getenv('HF_TOKEN')
model_id = 'meta-llama/Llama-2-70b-chat-hf'
text_gen = TextGenerator(token, model_id)
ui = GradioUI(text_gen)
ui.generate_interface()