Update app.py

app.py

@@ -1,14 +1,14 @@
-import streamlit as st
-import streamlit.components.v1 as components
+# app.py
+from flask import Flask, render_template, request, jsonify
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import json
-from streamlit.runtime.scriptrunner import add_script_run_ctx
-import threading
+from functools import lru_cache
 
-# تكوين النموذج والتوكنايزر
-@st.cache_resource
+app = Flask(__name__)
+
+@lru_cache(maxsize=1)
 def load_model():
+    """Load model and tokenizer with caching"""
     tokenizer = AutoTokenizer.from_pretrained("amd/AMD-OLMo-1B")
     model = AutoModelForCausalLM.from_pretrained(
         "amd/AMD-OLMo-1B",
@@ -18,8 +18,8 @@ def load_model():
     return model, tokenizer
 
 def generate_response(prompt, model, tokenizer):
+    """Generate response from the model"""
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -30,62 +30,29 @@ def generate_response(prompt, model, tokenizer):
             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()
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response.replace(prompt, "").strip()
 
-def main():
-    st.set_page_config(
-        page_title="سبيدي",
-        page_icon="💬",
-        layout="wide"
-    )
-    
-    # إخفاء عناصر Streamlit الافتراضية
-    hide_streamlit_style = """
-        <style>
-        #MainMenu {visibility: hidden;}
-        footer {visibility: hidden;}
-        header {visibility: hidden;}
-        </style>
-    """
-    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()
-    
-    # معالجة الرسائل الواردة من 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 message_handler(message_data):
-        ctx = add_script_run_ctx()
-        response = handle_message(message_data)
-        ctx.enqueue(json.dumps(response))
-    
-    # عرض الواجهة
-    components.html(
-        read_html(),
-        height=800,
-        on_message=message_handler
-    )
+@app.route('/')
+def home():
+    return render_template('index.html')
+
+@app.route('/message', methods=['POST'])
+def message():
+    try:
+        data = request.json
+        user_message = data.get('message', '')
+        
+        if not user_message:
+            return jsonify({"response": "عذراً، لم أفهم رسالتك"})
+        
+        model, tokenizer = load_model()
+        response = generate_response(user_message, model, tokenizer)
+        
+        return jsonify({"response": response})
+    
+    except Exception as e:
+        return jsonify({"response": f"عذراً، حدث خطأ: {str(e)}"})
 
-if __name__ == "__main__":
-    main()
+if __name__ == '__main__':
+    app.run(debug=True)