merterbak commited on
Commit
a4b21e5
·
verified ·
1 Parent(s): 24071a4

No streaming quick demo

Browse files
Files changed (1) hide show
  1. app.py +87 -0
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline, TextIteratorStreamer
2
+ import torch
3
+ from threading import Thread
4
+ import gradio as gr
5
+ import spaces
6
+
7
+ model_id = "openai/gpt-oss-20b"
8
+
9
+ pipe = pipeline(
10
+ "text-generation",
11
+ model=model_id,
12
+ torch_dtype="auto",
13
+ device_map="auto",
14
+ )
15
+
16
+ def format_conversation_history(chat_history):
17
+ messages = []
18
+ for item in chat_history:
19
+ role = item["role"]
20
+ content = item["content"]
21
+ if isinstance(content, list):
22
+ content = content[0]["text"] if content and "text" in content[0] else str(content)
23
+ messages.append({"role": role, "content": content})
24
+ return messages
25
+
26
+ @spaces.GPU(duration=120)
27
+ def generate_response(input_data, chat_history, max_new_tokens, system_prompt, temperature, top_p, top_k, repetition_penalty):
28
+ new_message = {"role": "user", "content": input_data}
29
+ system_message = [{"role": "system", "content": system_prompt}] if system_prompt else []
30
+ processed_history = format_conversation_history(chat_history)
31
+ messages = system_message + processed_history + [new_message]
32
+
33
+ streamer = TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, skip_special_tokens=True)
34
+ generation_kwargs = {
35
+ "streamer": streamer,
36
+ "max_new_tokens": max_new_tokens,
37
+ "do_sample": True,
38
+ "temperature": temperature,
39
+ "top_p": top_p,
40
+ "top_k": top_k,
41
+ "repetition_penalty": repetition_penalty
42
+ }
43
+ thread = Thread(target=pipe, args=(messages,), kwargs=generation_kwargs)
44
+ thread.start()
45
+ outputs = []
46
+ for text_chunk in streamer:
47
+ outputs.append(text_chunk)
48
+ yield "".join(outputs)
49
+
50
+ demo = gr.ChatInterface(
51
+ fn=generate_response,
52
+ additional_inputs=[
53
+ gr.Slider(label="Max new tokens", minimum=64, maximum=4096, step=1, value=1024),
54
+ gr.Textbox(
55
+ label="System Prompt",
56
+ value="You are a helpful assistant. Reasoning: medium",
57
+ lines=4,
58
+ placeholder="Change system prompt"
59
+ ),
60
+ gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7),
61
+ gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
62
+ gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=50),
63
+ gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, step=0.05, value=1.0)
64
+ ],
65
+ examples=[
66
+ [{"text": "What are the benefits of open weight AI models"}],
67
+ [{"text": "Write a Python function to calculate the Fibonacci sequence"}],
68
+ [{"text": "Explain Newton laws clearly and concisely."}],
69
+ ],
70
+ cache_examples=False,
71
+ type="messages",
72
+ description="""
73
+ # gpt-oss-20b
74
+ You can adjust reasoning level in the system prompt like "Reasoning: high".
75
+ """,
76
+ fill_height=True,
77
+ textbox=gr.Textbox(
78
+ label="Query Input",
79
+ placeholder="Type your prompt"
80
+ ),
81
+ stop_btn="Stop Generation",
82
+ multimodal=False,
83
+ theme=gr.themes.Soft()
84
+ )
85
+
86
+ if __name__ == "__main__":
87
+ demo.launch(share=True)