import gradio as gr
import galai as gal
import re
import urllib

model = gal.load_model("base")

def cite(prompt):
    text = model.generate(prompt+'[START_REF]')
    pattern = r'\[START_REF\](.*?)\[END_REF\]'
    references = re.findall(pattern, text)

    base_url = 'https://scholar.google.com/scholar?q='
    search_links = [base_url + urllib.parse.quote(reference) for reference in references]

    # Print the constructed links
    for i in range(len(references)):

        references[i] = f'<a href="{search_links[i]}" target="_blank">{references[i]}</a>'
    
    references = list(set(references))
    return '<br>'.join(references)

iface = gr.Interface(fn=cite, inputs="text", outputs="html",examples=["The cosine scheduler has been used in several papers as a scheduler for training large language models.",
	"We propose a new simple network architecture based on the original Transformer.",
	"The loss scales as a power-law with model size, dataset size, and the amount of compute used for training, with some trends spanning more than seven orders of magnitude.",
    "Molecular species that emerge and destroy during the birth of stars can be used to track the starforming processes within molecular clumps and cores",
    "Large Language Models (LLMs) have issues with document question answering (QA) in situations where the document is unable to fit in the small context length of an LLM"
	]) 
iface.launch()