import os import torch import torch.nn as nn import torch.nn.functional as F import sentencepiece as spm import streamlit as st import time import asyncio import sys from huggingface_hub import hf_hub_download os.environ["STREAMLIT_WATCHED_FILES"] = "" if sys.platform.startswith("win"): asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)] # # or simply: # torch.classes.__path__ = [] # Page configuration st.set_page_config( page_title="Mehfil-e-Sukhan", page_icon="📜", layout="wide" # Using wide layout for side-by-side content ) # Custom spinner implementation def custom_spinner(): spinner_html = """

Creating your poetry...

""" return spinner_html # Apply custom CSS for a professional UI with true vertical split st.markdown(""" """, unsafe_allow_html=True) # ------------------------------- # 1. Model Definition # ------------------------------- class BiLSTMLanguageModel(nn.Module): def __init__(self, vocab_size, embed_dim=512, hidden_dim=768, num_layers=3, dropout=0.2): super(BiLSTMLanguageModel, self).__init__() self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=0) self.lstm = nn.LSTM( input_size=embed_dim, hidden_size=hidden_dim, num_layers=num_layers, batch_first=True, bidirectional=True, dropout=dropout ) self.fc = nn.Linear(hidden_dim * 2, vocab_size) def forward(self, x, hidden=None): emb = self.embed(x) out, hidden = self.lstm(emb, hidden) logits = self.fc(out) return logits, hidden # ------------------------------- # 2. Poetry Generation Function # ------------------------------- @st.cache_resource def load_models(): # Define your model repository repo_id = "zaiffi/Mehfil-e-Sukhan" # Replace with your actual username/model-name # Download SentencePiece Model try: sp_model_path = hf_hub_download(repo_id=repo_id, filename="urdu_sp.model") sp = spm.SentencePieceProcessor() sp.load(sp_model_path) except Exception as e: st.error(f"Error loading SentencePiece model: {e}") return None, None # Initialize & Load the Trained Model vocab_size = sp.get_piece_size() model = BiLSTMLanguageModel(vocab_size, embed_dim=512, hidden_dim=768, num_layers=3, dropout=0.2) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) try: weights_path = hf_hub_download(repo_id=repo_id, filename="model_weights.pth") model.load_state_dict(torch.load(weights_path, map_location=device)) model.eval() except Exception as e: st.error(f"Error loading model weights: {e}") return sp, None return sp, model def generate_poetry_nucleus(model, sp, start_word, num_words=12, temperature=1.2, top_p=0.85): """ Generates poetry using nucleus (top-p) sampling. 'num_words' means 1 starting word + (num_words - 1) generated tokens. Output is formatted to 6 words per line. """ device = next(model.parameters()).device start_ids = sp.encode_as_ids(start_word) input_ids = [2] + start_ids # BOS token=2 input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device) with torch.no_grad(): logits, hidden = model(input_tensor) generated_ids = input_ids[:] for _ in range(num_words - 1): # generate one less token than num_words last_logits = logits[:, -1, :] scaled_logits = last_logits / temperature sorted_logits, sorted_indices = torch.sort(scaled_logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) filtered_indices = cumulative_probs > top_p if torch.all(filtered_indices): filtered_indices[-1] = False sorted_indices = sorted_indices[~filtered_indices] sorted_logits = sorted_logits[~filtered_indices] if len(sorted_indices) > 0: next_token_id = sorted_indices[torch.multinomial(F.softmax(sorted_logits, dim=-1), 1).item()].item() else: next_token_id = torch.argmax(last_logits).item() generated_ids.append(next_token_id) next_input = torch.tensor([[next_token_id]], dtype=torch.long, device=device) logits, hidden = model(next_input, hidden) # Decode & format with indentation for every second line generated_text = sp.decode_ids(generated_ids[1:]) # skip BOS words = generated_text.split() # Group words into lines of 6 words each lines = [] for i in range(0, len(words), 6): line = " ".join(words[i:i+6]) lines.append(line) # Format lines with indentation for even-numbered lines (0-indexed) formatted_lines = [] for i, line in enumerate(lines): if i % 2 == 1: # Every second line (1, 3, 5...) formatted_lines.append(f'

{line}

') else: formatted_lines.append(f'

{line}

') formatted_text = "\n".join(formatted_lines) return formatted_text # ------------------------------- # 3. Main Application # ------------------------------- def main(): # Load models sp, model = load_models() # Setup session state for storing poetry if 'poetry' not in st.session_state: st.session_state.poetry = "" # Define spinner_placeholder at the function level, before it's used spinner_placeholder = st.empty() # Define it here, at the main function level # Title and subtitle st.markdown('

Mehfil-e-Sukhan

', unsafe_allow_html=True) st.markdown('

Har Lafz Ek Mehfil

', unsafe_allow_html=True) # Create the true vertical split col1, col2 = st.columns([3, 7]) # 30% left, 70% right # Left Column (30%) - Control Settings with col1: st.markdown('

Control Settings

', unsafe_allow_html=True) # Number of Words Control st.markdown('

', unsafe_allow_html=True) st.markdown('

Number of Words

', unsafe_allow_html=True) st.markdown('

Total words in the generated poetry

', unsafe_allow_html=True) # Number of Words Control num_words = st.selectbox("Number of Words", options=[12, 18, 24, 30, 36, 42, 48], index=0, label_visibility="collapsed") st.markdown('

', unsafe_allow_html=True) # Creativity (Temperature) Control st.markdown('

', unsafe_allow_html=True) st.markdown('

Creativity

', unsafe_allow_html=True) st.markdown('

Higher values generate more unique poetry

', unsafe_allow_html=True) # Creativity (Temperature) Control temperature = st.slider("Creativity", 0.5, 2.0, 1.2, 0.1, key="creativity", label_visibility="collapsed") st.markdown('

', unsafe_allow_html=True) # Focus (Top-p) Control st.markdown('

', unsafe_allow_html=True) st.markdown('

Focus

', unsafe_allow_html=True) st.markdown('

Higher focus makes the AI stick to probable words

', unsafe_allow_html=True) # Focus (Top-p) Control top_p = st.slider("Focus", 0.5, 1.0, 0.85, 0.05, key="focus", label_visibility="collapsed") st.markdown('

', unsafe_allow_html=True) st.markdown('', unsafe_allow_html=True) # Right Column (70%) - Input and Output with col2: # Add a container with top margin st.markdown('

', unsafe_allow_html=True) # Starting Word Input with proper left margin st.markdown('

Starting Word/Phrase

', unsafe_allow_html=True) # Starting Word Input start_word = st.text_input("Starting Word/Phrase", value="ishq", placeholder="Enter a Roman Urdu word", label_visibility="collapsed") # Generate Button with custom width st.markdown('

', unsafe_allow_html=True) generate_button = st.button("Generate Poetry") st.markdown('

', unsafe_allow_html=True) # And then in your button click handler: if generate_button: if not sp or not model: st.error("Models not properly loaded. Please check the model files.") else: # Create a placeholder for the spinner spinner_placeholder = st.empty() # Show custom spinner spinner_placeholder.markdown(custom_spinner(), unsafe_allow_html=True) # Generate poetry time.sleep(0.1) # Add slight delay for smooth transition st.session_state.poetry = generate_poetry_nucleus( model, sp, start_word, num_words=num_words, temperature=temperature, top_p=top_p ) # Clear the spinner once poetry is generated spinner_placeholder.empty() # Display generated poetry if st.session_state.poetry: st.markdown(f"""

{st.session_state.poetry}

""", unsafe_allow_html=True) st.markdown('

', unsafe_allow_html=True) # Footer with quote st.markdown(""" """, unsafe_allow_html=True) if __name__ == "__main__": main()