Update app.py

app.py

@@ -1,10 +1,40 @@
 import streamlit as st
 import streamlit.components.v1 as components
-from transformers import pipeline
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import json
+from streamlit.runtime.scriptrunner import add_script_run_ctx
+import threading
+
+# تكوين النموذج والتوكنايزر
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("amd/AMD-OLMo-1B")
+    model = AutoModelForCausalLM.from_pretrained(
+        "amd/AMD-OLMo-1B",
+        torch_dtype=torch.float16,
+        device_map="auto"
+    )
+    return model, tokenizer
+
+def generate_response(prompt, model, tokenizer):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_length=200,
+            num_return_sequences=1,
+            temperature=0.7,
+            top_p=0.9,
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.replace(prompt, "").strip()
 
 def main():
-    # تكوين الصفحة
     st.set_page_config(
         page_title="سبيدي",
         page_icon="💬",
@@ -21,60 +51,41 @@ def main():
     """
     st.markdown(hide_streamlit_style, unsafe_allow_html=True)
     
+    # تحميل النموذج والتوكنايزر
+    model, tokenizer = load_model()
+    
     # قراءة ملف HTML
     def read_html():
         with open('index.html', 'r', encoding='utf-8') as file:
             return file.read()
     
-    # تهيئة نموذج Hugging Face
-    @st.cache_resource
-    def load_model():
-        # يمكنك تغيير اسم النموذج حسب احتياجك
-        return pipeline("text-generation", model="facebook/opt-350m", device=-1)
+    # معالجة الرسائل الواردة من JavaScript
+    def handle_message(message_data):
+        try:
+            data = json.loads(message_data)
+            user_message = data.get('message', '')
+            
+            if user_message:
+                response = generate_response(user_message, model, tokenizer)
+                return {"response": response}
+            
+            return {"response": "عذراً، لم أفهم رسالتك"}
+            
+        except Exception as e:
+            return {"response": f"عذراً، حدث خطأ: {str(e)}"}
     
-    # معالجة الرسائل
-    def process_message(message):
-        model = load_model()
-        response = model(message, max_length=100, num_return_sequences=1)
-        return response[0]['generated_text']
+    # تكوين معالج الرسائل
+    def message_handler(message_data):
+        ctx = add_script_run_ctx()
+        response = handle_message(message_data)
+        ctx.enqueue(json.dumps(response))
     
-    # إعداد API للتواصل بين JavaScript و Python
-    if 'messages' not in st.session_state:
-        st.session_state.messages = []
-    
-    # استقبال الرسائل من JavaScript
-    result = components.html(
-        read_html() + """
-        <script>
-        // تعديل دالة sendMessage لتتواصل مع Python
-        async function sendMessage() {
-            var input = document.getElementById('message-input');
-            var message = input.value.trim();
-            if (message !== '') {
-                appendMessage('user-message', message);
-                input.value = '';
-                
-                // إرسال الرسالة إلى Python
-                const response = await fetch('/_stcore/message', {
-                    method: 'POST',
-                    headers: {'Content-Type': 'application/json'},
-                    body: JSON.stringify({message: message})
-                });
-                
-                const data = await response.json();
-                appendMessage('bot-message', data.response);
-            }
-        }
-        </script>
-        """,
-        height=800
+    # عرض الواجهة
+    components.html(
+        read_html(),
+        height=800,
+        on_message=message_handler
     )
-    
-    # معالجة الرسائل المستلمة من JavaScript
-    if st.session_state.messages:
-        last_message = st.session_state.messages[-1]
-        response = process_message(last_message)
-        st.session_state.messages.append(response)
 
 if __name__ == "__main__":
     main()