import streamlit as st
from transformers import T5ForConditionalGeneration, T5Tokenizer
import torch
import spacy
import nltk
from b import b
#nltk.download('punkt')
from nltk.tokenize import sent_tokenize

# Load spaCy model
nlp = spacy.load("en_core_web_sm")

# Load T5 model and tokenizer
model_name = "DevBM/t5-large-squad"
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)

# Function to extract keywords using spaCy
def extract_keywords(text):
    doc = nlp(text)
    keywords = set()
    # Extract named entities
    for entity in doc.ents:
        keywords.add(entity.text)
    # Extract nouns and proper nouns
    for token in doc:
        if token.pos_ in ["NOUN", "PROPN"]:
            keywords.add(token.text)
    return list(keywords)

# Function to map keywords to sentences
def map_keywords_to_sentences(text, keywords):
    sentences = sent_tokenize(text)
    keyword_sentence_mapping = {}
    for keyword in keywords:
        for i, sentence in enumerate(sentences):
            if keyword in sentence:
                # Combine current sentence with surrounding sentences for context
                start = max(0, i-1)
                end = min(len(sentences), i+2)
                context = ' '.join(sentences[start:end])
                if keyword not in keyword_sentence_mapping:
                    keyword_sentence_mapping[keyword] = context
                else:
                    keyword_sentence_mapping[keyword] += ' ' + context
    return keyword_sentence_mapping

# Function to generate questions
def generate_question(context, answer):
    input_text = f"<context> {context} <answer> {answer}"
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    outputs = model.generate(input_ids)
    question = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return question

# Streamlit interface
st.title("Question Generator from Text")
text = st.text_area("Enter text here:")
if st.button("Generate Questions"):
    if text:
        keywords = extract_keywords(text)
        keyword_sentence_mapping = map_keywords_to_sentences(text, keywords)
        
        st.subheader("Generated Questions:")
        for keyword, context in keyword_sentence_mapping.items():
            question = generate_question(context, keyword)
            st.write(f"**Context:** {context}")
            st.write(f"**Answer:** {keyword}")
            st.write(f"**Question:** {question}")
            st.write("---")
    else:
        st.write("Please enter some text to generate questions.")