File size: 1,854 Bytes
c1811af
9831428
a381bc0
9831428
 
 
 
 
 
 
 
9522bb7
9831428
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9522bb7
 
9831428
 
 
4970856
 
 
a381bc0
 
4970856
a381bc0
4970856
a381bc0
4970856
 
 
 
 
 
 
9522bb7
 
4970856
 
 
 
 
9831428
c1811af
4970856
 
c1811af
4970856
 
 
c1811af
e537f35
4970856
c1811af
4970856
 
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
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import os
import gradio as gr
import wikipediaapi as wk
from transformers import (
    TokenClassificationPipeline,
    AutoModelForTokenClassification,
    AutoTokenizer,
)
from transformers.pipelines import AggregationStrategy
import numpy as np

# =====[ DEFINE PIPELINE ]===== #
class KeyphraseExtractionPipeline(TokenClassificationPipeline):
    def __init__(self, model, *args, **kwargs):
        super().__init__(
            model=AutoModelForTokenClassification.from_pretrained(model),
            tokenizer=AutoTokenizer.from_pretrained(model),
            *args,
            **kwargs
        )

    def postprocess(self, model_outputs):
        results = super().postprocess(
            model_outputs=model_outputs,
            aggregation_strategy=AggregationStrategy.SIMPLE,
        )
        return np.unique([result.get("word").strip() for result in results])

# =====[ LOAD PIPELINE ]===== #
model_name = "ml6team/keyphrase-extraction-kbir-inspec"
extractor = KeyphraseExtractionPipeline(model=model_name)

#TODO: add further preprocessing
def keyphrases_extraction(text: str) -> str:
    keyphrases = extractor(text)
    return keyphrases

def wikipedia_search(input: str) -> str:
    input = input.replace("\n", " ")
    keyphrases = keyphrases_extraction(input)
    wiki = wk.Wikipedia('en')
    
    try :
        #TODO: add better extraction and search
        page = wiki.page(keyphrases[0])
        return  page.summary
    except:
        return "I cannot answer this question"

# =====[ DEFINE INTERFACE ]===== #'
title = "Azza Chatbot"
examples = [
    ["Where is the Eiffel Tower?"],
    ["What is the population of France?"]
]


demo = gr.Interface(
    title = title,

    fn=wikipedia_search,
    inputs = "text", 
    outputs = "text",

    examples=examples,
    )

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