File size: 793 Bytes
45e64e0
f9588c9
5c3c196
f9588c9
 
5c3c196
f9588c9
 
 
45e64e0
5c3c196
f9588c9
 
 
 
 
 
 
 
 
 
 
45e64e0
 
f9588c9
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
import gradio as gr
import os
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForChatGPT

model_path = "cognitivecomputations/dolphin-2.7-mixtral-8x7b"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForChatGPT.from_pretrained(model_path)

@spaces.GPU
def chat(prompt):
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    output = model.generate(input_ids, max_length=1024, num_return_sequences=1, top_p=0.9, top_k=50, num_beams=2, early_stopping=True)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(value="Hello!", lines=5),
    outputs=gr.Textbox(label="Bot's Response", lines=5)
)

if __name__ == "__main__":
    demo.launch()