Create app.py

app.py

@@ -0,0 +1,117 @@
+import gradio as gr
+import torch
+import itertools
+import pandas as pd
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+
+
+model_name = 'philipp-zettl/t5-small-long-qa'
+qa_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+model_name = 'philipp-zettl/t5-small-qg'
+qg_model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained('google/flan-t5-small')
+
+# Move only the student model to GPU if available
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+qa_model = qa_model.to(device)
+qg_model = qg_model.to(device)
+
+max_questions = 1
+max_answers = 1
+
+
+def run_model(inputs, tokenizer, model, temperature=0.5, num_return_sequences=1):
+    all_outputs = []
+    for input_text in inputs:
+        model_inputs = tokenizer([input_text], max_length=512, padding=True, truncation=True)
+        input_ids = torch.tensor(model_inputs['input_ids']).to(device)
+        for sample in input_ids:
+            sample_outputs = []
+            with torch.no_grad():
+                sample_output = model.generate(
+                    input_ids[:1],
+                    max_length=85,
+                    temperature=temperature,
+                    do_sample=True,
+                    num_return_sequences=num_return_sequences,
+                    low_memory=True,
+                    num_beams=max(2, num_return_sequences),
+                    use_cache=True,
+                )
+                for i, sample_output in enumerate(sample_output):
+                    sample_output = sample_output.unsqueeze(0)
+                    sample_output = tokenizer.decode(sample_output[0], skip_special_tokens=True)
+                    sample_outputs.append(sample_output)
+
+            all_outputs.append(sample_outputs)
+    return all_outputs
+
+
+def gen(content, temperature_qg=0.5, temperature_qa=0.75, num_return_sequences_qg=1, num_return_sequences_qa=1):
+    inputs = [
+        f'context: {content}'
+    ]
+    question = run_model(inputs, tokenizer, qg_model, temperature_qg, num_return_sequences_qg)
+
+    inputs = list(itertools.chain.from_iterable([
+        [f'question: {q} {inputs[0]}' for q in q_set] for q_set in question
+    ]))
+    answer = run_model(inputs, tokenizer, qa_model, temperature_qa, num_return_sequences_qa)
+
+    questions = list(itertools.chain.from_iterable(question))
+    answers = list(itertools.chain.from_iterable(answer))
+
+    results = []
+    for idx, ans in enumerate(answers):
+        results.append({'question': questions[idx % num_return_sequences_qg], 'answer': ans})
+    return results
+
+
+def variable_outputs(k, max_elems=10):
+    k = int(k)
+    return [gr.Text(visible=True)] * k + [gr.Text(visible=False)] * (max(max_elems, 10)- k)
+
+
+def set_outputs(content, max_elems=10):
+    c = eval(content)
+    print('received content: ', c)
+    return [gr.Text(value=t, visible=True) for t in c] + [gr.Text(visible=False)] * (max(max_elems, 10) - len(c))
+
+
+def create_file_download(qnas):
+    with open('qnas.tsv', 'w') as f:
+        for idx, qna in qnas.iterrows():
+            f.write(qna['Question'] + '\t' + qna['Answer'])
+            if idx < len(qnas) - 1:
+                f.write('\n')
+    return 'qnas.tsv'
+
+
+with gr.Blocks() as demo:
+    with gr.Row(equal_height=True):
+        with gr.Group("Content"):
+            content = gr.Textbox(label='Content', lines=15, placeholder='Enter text here', max_lines=10_000)
+        with gr.Group("Settings"):
+            temperature_qg = gr.Slider(label='Temperature QG', value=0.5, minimum=0, maximum=1, step=0.01)
+            temperature_qa = gr.Slider(label='Temperature QA', value=0.75, minimum=0, maximum=1, step=0.01)
+            num_return_sequences_qg = gr.Number(label='Number Questions', value=max_questions, minimum=1, step=1, maximum=max(max_questions, 10))
+            num_return_sequences_qa = gr.Number(label="Number Answers", value=max_answers, minimum=1, step=1, maximum=max(max_questions, 10))
+
+    with gr.Row():
+        gen_btn = gr.Button("Generate")
+
+    @gr.render(inputs=[content, temperature_qg, temperature_qa, num_return_sequences_qg, num_return_sequences_qa], triggers=[gen_btn.click])
+    def render_results(content, temperature_qg, temperature_qa, num_return_sequences_qg, num_return_sequences_qa):
+        qnas = gen(content, temperature_qg, temperature_qa, num_return_sequences_qg, num_return_sequences_qa)
+        df = gr.Dataframe(
+            value=[u.values() for u in qnas],
+            headers=['Question', 'Answer'],
+            col_count=2,
+            wrap=True
+        )
+        pd_df = pd.DataFrame([u.values() for u in qnas], columns=['Question', 'Answer'])
+
+        download = gr.DownloadButton(label='Download (without headers)', value=create_file_download(pd_df))
+
+
+demo.launch()