app.py

import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM  # Fixed import name
import gradio as gr
import warnings
warnings.filterwarnings('ignore')

class SpaceBot:
    def __init__(self):
        # Initialize model and tokenizer
        print("Loading model and tokenizer...")
        self.model_name = "google/flan-t5-base"  # Using a more stable model
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name)
        
        # Space facts database
        self.space_facts = {
            "apollo": [
                "The Apollo program ran from 1961 to 1972.",
                "Apollo 11 was the first mission to land humans on the Moon.",
                "Neil Armstrong and Buzz Aldrin walked on the Moon on July 20, 1969.",
                "The Apollo program included 17 space missions in total."
            ],
            "iss": [
                "The International Space Station orbits Earth at about 400 kilometers above the surface.",
                "The ISS has been continuously occupied since November 2000.",
                "The ISS travels at approximately 7.66 kilometers per second.",
                "The space station is about the size of a football field."
            ],
            "mars": [
                "Mars is often called the Red Planet due to its reddish appearance.",
                "NASA's Perseverance rover landed on Mars in 2021.",
                "Mars has two small moons: Phobos and Deimos.",
                "The first successful Mars landing was by NASA's Viking 1 in 1976."
            ],
            "telescope": [
                "The Hubble Space Telescope was launched in 1990.",
                "The James Webb Space Telescope is the largest space telescope ever built.",
                "Space telescopes can see more clearly than ground-based telescopes.",
                "The Hubble has made over 1.5 million observations."
            ],
            "general": [
                "The first human in space was Yuri Gagarin in 1961.",
                "SpaceX achieved the first successful landing of a reusable rocket in 2015.",
                "There are currently over 5,000 satellites orbiting Earth.",
                "Light takes about 8 minutes to travel from the Sun to Earth."
            ]
        }
    
    def get_relevant_facts(self, query):
        """Get relevant facts based on keywords in the query"""
        query = query.lower()
        relevant_facts = []
        
        # Check each category for relevant facts
        for category, facts in self.space_facts.items():
            if category in query or any(keyword in query for keyword in category.split()):
                relevant_facts.extend(facts)
        
        # If no specific facts found, return general facts
        if not relevant_facts:
            relevant_facts = self.space_facts["general"]
        
        return " ".join(relevant_facts[:2])  # Return two relevant facts

    def generate_response(self, query):
        try:
            # Get relevant facts as context
            context = self.get_relevant_facts(query)
            
            # Prepare input text
            input_text = f"Using this context: {context}\nAnswer this question: {query}"
            
            # Tokenize input
            inputs = self.tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
            
            # Generate response
            outputs = self.model.generate(
                inputs.input_ids,
                max_length=150,
                min_length=30,
                num_beams=4,
                temperature=0.7,
                no_repeat_ngram_size=2
            )
            
            # Decode response
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            return response
        except Exception as e:
            print(f"Error generating response: {str(e)}")
            return "I apologize, but I couldn't generate a proper response. Please try asking your question in a different way."

def chat_with_bot(message, history):
    """Handle chat interactions"""
    try:
        response = bot.generate_response(message)
        return [(response, None)]
    except Exception as e:
        print(f"Error in chat interface: {str(e)}")
        return [("I encountered an error. Please try asking your question again.", None)]

print("Initializing SpaceBot...")
# Initialize the bot
bot = SpaceBot()
print("SpaceBot initialization complete!")

# Create Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🚀 Space Explorer Chat
    Ask me anything about space exploration and astronomy!
    """)
    
    chatbot = gr.ChatInterface(
        chat_with_bot,
        chatbot=gr.Chatbot(height=400),
        textbox=gr.Textbox(
            placeholder="Ask about space missions, planets, or astronomy...",
            container=False
        ),
        title="Space Explorer",
        description="Your guide to space exploration and astronomy",
        theme="soft",
        examples=[
            "What was the Apollo program?",
            "Tell me about the International Space Station",
            "What do you know about Mars exploration?",
            "Tell me about space telescopes",
            "What are some interesting space facts?"
        ],
    )
    
    gr.Markdown("""
    ### Topics I can help with:
    - Space missions and history
    - Planets and astronomy
    - Space stations and satellites
    - Space telescopes
    - General space facts
    """)

# Launch the interface
demo.launch(share=True)