import streamlit as st
import requests
from transformers import AutoTokenizer, AutoModel

# Load SciBERT pre-trained model and tokenizer
model_name = "allenai/scibert_scivocab_uncased"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def calculate_similarity(claim, document):
    # Tokenize claim and document
    inputs = tokenizer.encode_plus(claim, document, return_tensors='pt', padding=True, truncation=True)
    
    # Generate embeddings for claim and document
    with torch.no_grad():
        claim_embeddings = model(**inputs)['pooler_output']
    
    # Compute cosine similarity between embeddings
    similarity = torch.cosine_similarity(claim_embeddings, document_embeddings).item()
    
    return similarity

def search_papers(user_input):
    # Implement your code to fetch search results from the desired source (e.g., arXiv, Semantic Scholar)
    # ...

    # For the purpose of this example, we'll use dummy data
    search_results = [
        {
            'title': 'Paper 1 Title',
            'abstract': 'Paper 1 Abstract',
            'authors': ['Author 1', 'Author 2'],
            'url': 'https://example.com/paper1'
        },
        {
            'title': 'Paper 2 Title',
            'abstract': 'Paper 2 Abstract',
            'authors': ['Author 3', 'Author 4'],
            'url': 'https://example.com/paper2'
        },
        {
            'title': 'Paper 3 Title',
            'abstract': 'Paper 3 Abstract',
            'authors': ['Author 5', 'Author 6'],
            'url': 'https://example.com/paper3'
        }
    ]

    return search_results

st.title('The Substantiator')

user_input = st.text_input('Input your claim')

if st.button('Substantiate'):
    search_results = search_papers(user_input)
    if search_results is not None and len(search_results) > 0:
        for result in search_results:
            st.write(result["title"])
            st.write(result["abstract"])
            st.write("Authors: ", ", ".join(result["authors"]))
            similarity = calculate_similarity(user_input, result["abstract"])
            st.write("Similarity Score: ", similarity)
            st.write("-----")
    else:
        st.write("No results found.")