from flask import Flask, request, jsonify
import pymongo
import requests
import json
from transformers import MT5ForConditionalGeneration, MT5Tokenizer
import torch
from torch.nn.functional import softmax

app = Flask(__name__)

# MongoDB bağlantısı
mongo_client = pymongo.MongoClient('mongodb://localhost:27017/')
db = mongo_client['EgitimDatabase']
collection = db['test']

# Model ve tokenizer'ı yükleme
model_name = "alan-turing-institute/mt5-large-finetuned-mnli-xtreme-xnli"
tokenizer = MT5Tokenizer.from_pretrained(model_name)
model = MT5ForConditionalGeneration.from_pretrained(model_name)
model.eval()

def process_nli(premise: str, hypothesis: str):
    """NLI formatına uygun hale getirir"""
    return "".join(['xnli: premise: ', premise, ' hypothesis: ', hypothesis])

def generate_text(title, keywords, subheadings):
    # Prompt oluşturma
    prompt = f"Başlık: {title}\nAnahtar Kelimeler: {keywords}\nAlt Başlıklar: {subheadings}\n\nBu bilgilerle bir metin oluşturun:"
    inputs = tokenizer(prompt, return_tensors="pt")

    # Modelden çıktı alma
    outputs = model.generate(**inputs, output_scores=True, return_dict_in_generate=True, num_beams=5)
    generated_text = tokenizer.decode(outputs.sequences[0], skip_special_tokens=True)
    
    return generated_text

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    title = data.get('title')
    keywords = data.get('keywords')
    subheadings = data.get('subheadings')
    
    # MongoDB'den ilgili verileri çekme
    query = {
        'title': title,
        'keywords': {'$in': keywords.split(',')},
        'subheadings': {'$in': subheadings.split(',')}
    }
    documents = list(collection.find(query))

    if not documents:
        return jsonify({'error': 'No documents found matching the query'}), 404

    # Verilerle metin oluşturma
    generated_texts = []
    for doc in documents:
        generated_text = generate_text(doc['title'], doc['keywords'], doc['subheadings'])
        generated_texts.append(generated_text)

    # Sonuçları döndürme
    response = {
        'title': title,
        'keywords': keywords,
        'subheadings': subheadings,
        'generated_texts': generated_texts
    }
    
    return jsonify(response)

if __name__ == "__main__":
    app.run(host='0.0.0.0', port=8080)