import gradio
import os
import time
import csv
import datetime
from transformers import RobertaTokenizer, T5ForConditionalGeneration

def evaluate(sentence):
    tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-base')
    model = T5ForConditionalGeneration.from_pretrained('Salesforce/codet5-base-multi-sum')

    # Prepare the input text
    input_text = sentence.strip()
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    # Generate a summary
    generated_ids = model.generate(input_ids, max_length=20)
    summary = tokenizer.decode(generated_ids[0], skip_special_tokens=True)

    return summary

def predict(sentence):    
    timestamp = datetime.datetime.now().isoformat()    
    start_time = time.time()    
    predictions = evaluate(sentence)    
    elapsed_time = time.time() - start_time      
    output = predictions
    print(f"Sentence: {sentence} \nPrediction: {predictions}")    
    return output
    
gradio.Interface(
    fn=predict,
    inputs=gradio.inputs.Textbox(label="Enter Code Snippet:", placeholder="Type here...", lines=2),
    outputs="text",
    allow_flagging='never',
    title="Code Summarization"
).launch()