Update app.py

app.py

@@ -1,46 +1,127 @@
-import requests
 import streamlit as st
-import random
 import time
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-st.title("Rasa Chatbot Demo")
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Display chat messages from history on app rerun
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Accept user input
-if user_input := st.chat_input("Ask me a question, for example: when can I expect the test results? is the test good for kids?"):
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": user_input})
-    # Display user message in chat message container
-    with st.chat_message("user"):
-        st.markdown(user_input)
-
-    # Send user input to Rasa webhook
-    payload = {"sender": "user", "message": user_input}
-    response = requests.post('https://heyal-chatbot.hf.space/webhooks/rest/webhook', json=payload)
-    bot_reply = response.json()
-
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        message_placeholder = st.empty()
+class LLAMAChatbot:
+    def __init__(self):
+        st.title("LLAMA Chatbot")
+        self.initialize_model()
+        self.initialize_session_state()
+        
+    def initialize_model(self):
+        """Initialize the LLAMA model and tokenizer"""
+        try:
+            @st.cache_resource
+            def load_model():
+                tokenizer = AutoTokenizer.from_pretrained("joermd/llma-speedy")
+                model = AutoModelForCausalLM.from_pretrained(
+                    "joermd/llma-speedy",
+                    torch_dtype=torch.float16,
+                    device_map="auto"
+                )
+                return model, tokenizer
+            
+            self.model, self.tokenizer = load_model()
+            st.success("تم تحميل النموذج بنجاح!")
+        except Exception as e:
+            st.error(f"حدث خطأ أثناء تحميل النموذج: {str(e)}")
+            st.stop()
+            
+    def initialize_session_state(self):
+        """Initialize chat history if it doesn't exist"""
+        if "messages" not in st.session_state:
+            st.session_state.messages = []
+            
+    def display_chat_history(self):
+        """Display all messages from chat history"""
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+                
+    def add_message(self, role, content):
+        """Add a message to the chat history"""
+        st.session_state.messages.append({
+            "role": role,
+            "content": content
+        })
+        
+    def generate_response(self, user_input, max_length=1000):
+        """Generate response using LLAMA model"""
+        try:
+            # Prepare the input context with chat history
+            context = ""
+            for message in st.session_state.messages[-4:]:  # Use last 4 messages for context
+                if message["role"] == "user":
+                    context += f"Human: {message['content']}\n"
+                else:
+                    context += f"Assistant: {message['content']}\n"
+            
+            context += f"Human: {user_input}\nAssistant:"
+            
+            # Tokenize input
+            inputs = self.tokenizer(context, return_tensors="pt", truncation=True)
+            inputs = inputs.to(self.model.device)
+            
+            # Generate response
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    inputs["input_ids"],
+                    max_length=max_length,
+                    num_return_sequences=1,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.eos_token_id
+                )
+            
+            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # Extract only the assistant's response
+            response = response.split("Assistant:")[-1].strip()
+            
+            return response
+            
+        except Exception as e:
+            return f"عذراً، حدث خطأ أثناء توليد الإجابة: {str(e)}"
+            
+    def simulate_typing(self, message_placeholder, response):
+        """Simulate typing effect for bot response"""
         full_response = ""
-        if not bot_reply:
-            assistant_response = "I didn't understand your question. Could you please rephease?"
-        else:
-            assistant_response = random.choice(bot_reply)["text"]
-        # Simulate stream of response with milliseconds delay
-        for chunk in assistant_response.split():
+        for chunk in response.split():
             full_response += chunk + " "
             time.sleep(0.05)
-            # Add a blinking cursor to simulate typing
             message_placeholder.markdown(full_response + "▌")
         message_placeholder.markdown(full_response)
-    # Add assistant response to chat history
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+        return full_response
+        
+    def run(self):
+        """Main application loop"""
+        # Display existing chat history
+        self.display_chat_history()
+        
+        # Handle user input
+        if user_input := st.chat_input("اكتب رسالتك هنا..."):
+            # Display and save user message
+            self.add_message("user", user_input)
+            with st.chat_message("user"):
+                st.markdown(user_input)
+            
+            # Generate and display response
+            with st.chat_message("assistant"):
+                message_placeholder = st.empty()
+                with st.spinner("جاري التفكير..."):
+                    assistant_response = self.generate_response(user_input)
+                full_response = self.simulate_typing(message_placeholder, assistant_response)
+                self.add_message("assistant", full_response)
+
+if __name__ == "__main__":
+    # Set page config
+    st.set_page_config(
+        page_title="LLAMA Chatbot",
+        page_icon="🤖",
+        layout="wide"
+    )
+    
+    # Initialize and run the chatbot
+    chatbot = LLAMAChatbot()
+    chatbot.run()