Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import DistilBertTokenizer, DistilBertForQuestionAnswering
 import torch
 
 # Set page configuration
@@ -7,47 +7,48 @@ st.set_page_config(page_title="Urdu AI Chatbot", page_icon="🤖")
 
 # Title and description
 st.title("اردو AI چیٹ بوٹ")
-st.write("یہ ایک سادہ چیٹ بوٹ ہے جو اردو زبان میں جواب دیتا ہے۔ اپنا سوال درج کریں اور جواب حاصل کریں!")
+st.write("یہ ایک سادہ چیٹ بوٹ ہے جو اردو زبان میں جواب دیتا ہے۔ اپنا سوال درج کریں!")
 
-# Load the model and tokenizer from Hugging Face
+# Load a lightweight multilingual model
 @st.cache_resource
 def load_model():
     try:
-        model_name = "facebook/mbart-large-50"
-        model = MBartForConditionalGeneration.from_pretrained(model_name)
-        tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
-        # Set the target language to Urdu
-        tokenizer.tgt_lang = "ur_PK"
+        tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-multilingual-cased")
+        model = DistilBertForQuestionAnswering.from_pretrained("distilbert-base-multilingual-cased")
         return model, tokenizer
     except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
+        st.error(f"ماڈل لوڈ کرنے میں خرابی: {str(e)}")
         return None, None
 
 model, tokenizer = load_model()
 if model is None or tokenizer is None:
-    st.error("Model or tokenizer failed to load. Please check the logs or try again later.")
+    st.error("ماڈل یا ٹوکنائزر لوڈ نہیں ہوا۔ براہ کرم لاگز چیک کریں۔")
     st.stop()
 
-# Function to generate response
+# Function to generate a simple response
 def generate_response(user_input):
     try:
-        # Tokenize input and generate response
-        inputs = tokenizer(user_input, return_tensors="pt", padding=True)
-        # Force the output to be in Urdu
-        output_ids = model.generate(
-            **inputs,
-            max_length=50,
-            num_beams=5,
-            forced_bos_token_id=tokenizer.lang_code_to_id["ur_PK"]
-        )
-        response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        # Encode the input
+        inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
+        outputs = model(**inputs)
+        # Since this is a QA model, we'll simulate a response
+        start_scores = outputs.start_logits
+        end_scores = outputs.end_logits
+        start_idx = torch.argmax(start_scores)
+        end_idx = torch.argmax(end_scores) + 1
+        answer_tokens = inputs["input_ids"][0][start_idx:end_idx]
+        response = tokenizer.decode(answer_tokens, skip_special_tokens=True)
+        
+        # If the response is empty or too short, provide a fallback
+        if not response or len(response) < 3:
+            return "میں سمجھ گیا، لیکن براہ کرم مزید وضاحت کریں۔"
         return response
     except Exception as e:
         return f"معذرت، کچھ غلطی ہوئی: {str(e)}"
 
 # Chat interface
 if "messages" not in st.session_state:
-    st.session_state.messages = [{"role": "assistant", "content": "سلام! میں آپ کی مدد کیسے کر سکتا ہوں؟"}]
+    st.session_state.messages = [{"role": "assistant", "content": "سلام! میں آپ کی مدد کے لیے حاضر ہوں۔"}]
 
 # Display chat history
 for message in st.session_state.messages:
@@ -56,19 +57,16 @@ for message in st.session_state.messages:
 
 # User input
 if prompt := st.chat_input("اپنا سوال یہاں لکھیں..."):
-    # Add user message to chat history
     st.session_state.messages.append({"role": "user", "content": prompt})
     with st.chat_message("user"):
         st.write(prompt)
 
-    # Generate and display assistant response
     with st.chat_message("assistant"):
-        with st.spinner("غور کر رہا ہوں..."):
+        with st.spinner("جواب تیار کر رہا ہوں..."):
             response = generate_response(prompt)
         st.write(response)
         st.session_state.messages.append({"role": "assistant", "content": response})
 
-# Sidebar with additional info
+# Sidebar
 st.sidebar.title("معلومات")
-st.sidebar.write("یہ چیٹ بوٹ Hugging Face کے `facebook/mbart-large-50` ماڈل کا استعمال کرتا ہے۔")
-st.sidebar.write("مزید ماڈلز کے لیے [Hugging Face](https://huggingface.co/models) دیکھیں۔")
+st.sidebar.write("یہ `distilbert-base-multilingual-cased` ماڈل استعمال کرتا ہے۔")