farhananis005/LLM_Finetuning
1
1import gradio as gr
2import random
3from threading import Thread
4from queue import Queue
5
6# Import our new modules
7import config
8import backend
9
10# --- HELPER FUNCTIONS (Unchanged) ---
11def get_random_question(domain):
12 data_conf = config.DATASET_CONFIG[domain]
13 dataset = data_conf["dataset"]
14
15 if not dataset:
16 return "Failed to load dataset.", "N/A"
17
18 random_index = random.randint(0, len(dataset) - 1)
19 sample = dataset[random_index]
20
21 if domain == "Math":
22 question = sample[data_conf["question_col"]]
23 answer = sample[data_conf["answer_col"]]
24 elif domain == "Bio":
25 instruction = sample[data_conf["instruction_col"]]
26 bio_input = sample[data_conf["input_col"]]
27 answer = sample[data_conf["answer_col"]]
28 if bio_input and bio_input.strip():
29 question = f"**Instruction:**\n{instruction}\n\n**Input:**\n{bio_input}"
30 else:
31 question = instruction
32
33 return question, answer
34
35def update_domain_settings(domain):
36 models = list(config.ALL_MODELS[domain].keys())
37 def_base = next((m for m in models if "Base" in m), models[0])
38 def_ft = next((m for m in models if "Finetuned" in m), models[0])
39
40 q, a = get_random_question(domain)
41 return [
42 gr.Dropdown(choices=models, value=def_base),
43 gr.Dropdown(choices=models, value=def_ft),
44 gr.Textbox(value=q),
45 a,
46 gr.Markdown(visible=False)
47 ]
48
49def load_next_question(domain):
50 q, a = get_random_question(domain)
51 return [gr.Textbox(value=q), a, gr.Markdown(visible=False, value="")]
52
53def reveal_answer(hidden_answer):
54 return gr.Markdown(value=f"**Ground Truth Answer:**\n\n{hidden_answer}", visible=True)
55
56# --- CORE LOGIC (REBUILT FOR TRUE PARALLEL STREAMING) ---
57
58def stream_to_queue(model_id, prompt, lane, queue, key):
59 """
60 A worker function that runs in a thread.
61 It calls the streaming API and puts tokens into the queue.
62 """
63 try:
64 # call_modal_api is a generator
65 for token in backend.call_modal_api(model_id, prompt, lane):
66 queue.put((key, token))
67 except Exception as e:
68 queue.put((key, f"\n\nTHREAD ERROR: {e}"))
69 finally:
70 # When the stream is done, put a 'None' sentinel
71 queue.put((key, None))
72
73def run_comparison(domain, question, model_1_name, model_2_name):
74 # 1. Get IDs
75 id_1 = config.ALL_MODELS[domain].get(model_1_name)
76 id_2 = config.ALL_MODELS[domain].get(model_2_name)
77
78 # 2. Ask the Smart Router
79 lane_for_m1, lane_for_m2 = backend.router.get_routing_plan(id_1, id_2)
80
81 # 3. Create the Queue and Threads
82 q = Queue()
83
84 Thread(
85 target=stream_to_queue,
86 args=(id_1, question, lane_for_m1, q, 'm1')
87 ).start()
88
89 Thread(
90 target=stream_to_queue,
91 args=(id_2, question, lane_for_m2, q, 'm2')
92 ).start()
93
94 # 4. Listen to the Queue
95 text1 = ""
96 text2 = ""
97 m1_done = False
98 m2_done = False
99
100 # Clear boxes and start
101 yield "", "", gr.Markdown(visible=False)
102
103 while not (m1_done and m2_done):
104 # Wait for the next token from *either* thread
105 try:
106 key, token = q.get()
107 except Exception as e:
108 # This should ideally not happen
109 print(f"Queue error: {e}")
110 continue
111
112 # Check for the 'None' sentinel
113 if token is None:
114 if key == 'm1':
115 m1_done = True
116 elif key == 'm2':
117 m2_done = True
118 else:
119 # Append the new token
120 if key == 'm1':
121 text1 += token
122 elif key == 'm2':
123 text2 += token
124
125 # Yield the updated full text
126 yield text1, text2, gr.Markdown(visible=False)
127
128
129# --- UI BUILD (Unchanged) ---
130initial_question, initial_answer = get_random_question("Math")
131
132with gr.Blocks(theme=gr.themes.Soft()) as demo:
133 gr.Markdown(
134 """
135 # 🔬 LLM Finetuning Arena
136 ### Comparing Finetuned vs. Base Models on Specialized Tasks
137 """
138 )
139
140 hidden_answer_state = gr.State(value=initial_answer)
141
142 with gr.Row():
143 domain_radio = gr.Radio(
144 ["Math", "Bio"], label="1. Select Domain", value="Math"
145 )
146
147 with gr.Row():
148 question_box = gr.Textbox(
149 label="2. Question Prompt (Editable)",
150 value=initial_question, lines=5, scale=4
151 )
152 next_btn = gr.Button("Load Random Question 🔄", scale=1, min_width=100)
153
154 with gr.Row():
155 model_1_dd = gr.Dropdown(
156 label="3. Select Model 1 (Left)",
157 choices=list(config.ALL_MODELS["Math"].keys()),
158 value=next((m for m in config.ALL_MODELS["Math"] if "Base" in m))
159 )
160 model_2_dd = gr.Dropdown(
161 label="4. Select Model 2 (Right)",
162 choices=list(config.ALL_MODELS["Math"].keys()),
163 value=next((m for m in config.ALL_MODELS["Math"] if "Finetuned" in m))
164 )
165
166 with gr.Row():
167 run_btn = gr.Button("🚀 Run Comparison", variant="primary", scale=3)
168 show_answer_btn = gr.Button("Show Ground Truth Answer", scale=1)
169
170 answer_display_box = gr.Markdown(label="Ground Truth Answer", visible=False)
171
172 gr.Markdown("---")
173
174 with gr.Row():
175 output_1_box = gr.Markdown(label="Output: Model 1")
176 output_2_box = gr.Markdown(label="Output: Model 2")
177
178 # --- EVENTS (Unchanged) ---
179 domain_radio.change(
180 fn=update_domain_settings,
181 inputs=[domain_radio],
182 outputs=[model_1_dd, model_2_dd, question_box, hidden_answer_state, answer_display_box]
183 )
184
185 next_btn.click(
186 fn=load_next_question,
187 inputs=[domain_radio],
188 outputs=[question_box, hidden_answer_state, answer_display_box]
189 )
190
191 show_answer_btn.click(
192 fn=reveal_answer,
193 inputs=[hidden_answer_state],
194 outputs=[answer_display_box]
195 )
196
197 run_btn.click(
198 fn=run_comparison,
199 inputs=[domain_radio, question_box, model_1_dd, model_2_dd],
200 outputs=[output_1_box, output_2_box, answer_display_box]
201 )
202
203if __name__ == "__main__":
204 if not config.MY_AUTH_TOKEN:
205 print("⚠️ WARNING: ARENA_AUTH_TOKEN is not set.")
206 demo.launch()