Update app.py

app.py

@@ -7,7 +7,6 @@ import gradio as gr
 import asyncio
 from collections import OrderedDict
 from datetime import datetime
-from transformers import AutoProcessor, AutoTokenizer
 import requests
 from openai import OpenAI
 from telethon import TelegramClient, events
@@ -16,14 +15,12 @@ from huggingface_hub import InferenceClient
 import pymongo
 from pymongo import MongoClient
 
-# Load system prompt from a file
 def load_system_prompt():
     with open('prompt.txt', 'r') as file:
         return file.read()
 
 system_prompt = load_system_prompt()
 
-# Environment variables
 api_id = os.getenv('api_id')
 api_hash = os.getenv('api_hash')
 bot_token = os.getenv('bot_token')
@@ -35,24 +32,22 @@ model1 = os.getenv('model1')
 model2 = os.getenv('model2')
 mongoURI = os.getenv('MONGO_URI')
 
-# Initialize OpenAI and MongoDB clients
 openai_client = OpenAI(api_key=openai_api_key, base_url=api_url)
 mongo_client = MongoClient(mongoURI)
 db = mongo_client['Scarlett']
 chat_collection = db['chats']
 
-# Initialize Hugging Face models for image processing
+idefics_processor = AutoProcessor.from_pretrained(model1)
+idefics_client = InferenceClient(model2)
+tokenizer = AutoTokenizer.from_pretrained(model1)
 
-# Local cache for up to 5 users
 local_chat_history = OrderedDict()
 MAX_LOCAL_USERS = 5
 
-# Retrieve chat history from MongoDB
 def get_history_from_mongo(user_id):
     result = chat_collection.find_one({"user_id": user_id})
     return result.get("messages", []) if result else []
 
-# Store message in MongoDB (limit to last 99 messages)
 def store_message_in_mongo(user_id, role, content):
     chat_collection.update_one(
         {"user_id": user_id},
@@ -67,35 +62,25 @@ def store_message_in_mongo(user_id, role, content):
         upsert=True
     )
 
-# Get chat history from local cache or MongoDB
 def get_chat_history(user_id):
     if user_id in local_chat_history:
-        local_chat_history.move_to_end(user_id)  # Mark as most recently used
+        local_chat_history.move_to_end(user_id)
         return local_chat_history[user_id]
-
-    # Load from MongoDB if not in local cache
     history = get_history_from_mongo(user_id)
     local_chat_history[user_id] = history
-
     if len(local_chat_history) > MAX_LOCAL_USERS:
-        local_chat_history.popitem(last=False)  # Remove LRU user
-    
+        local_chat_history.popitem(last=False)
     return history
 
-# Update chat history (both local and MongoDB)
 def update_chat_history(user_id, role, content):
     if user_id not in local_chat_history:
         local_chat_history[user_id] = get_history_from_mongo(user_id)
-
     local_chat_history[user_id].append({"role": role, "content": content})
     local_chat_history.move_to_end(user_id)
-
     if len(local_chat_history) > MAX_LOCAL_USERS:
         local_chat_history.popitem(last=False)
-
     store_message_in_mongo(user_id, role, content)
 
-# Encode image to base64
 def encode_local_image(image_path):
     im = Image.fromarray(image_path)
     buffered = BytesIO()
@@ -104,18 +89,13 @@ def encode_local_image(image_path):
     image_base64 = base64.b64encode(image_bytes).decode('ascii')
     return image_base64
 
-# Describe image using the model with error handling
 def inference_calling_idefics(image_path, question=""):
     system_prompt = os.getenv('USER_PROMPT')
-    model_id = "HuggingFaceM4/idefics2-8b-chatty"
+    model_id = model2
     client = InferenceClient(model=model_id)
-    image_base64 = base64_encoded_image(image_path)
+    image_base64 = describe_image(image_path)
     image_info = f"data:image/png;base64,{image_base64}"
-    
-    # Include the system prompt before the user question
     prompt = f"{system_prompt}\n![]({image_info})\n{question}\n\n"
-    
-    # Adding do_sample=True to ensure different responses and max_tokens
     response = client.text_generation(
         prompt, 
         max_new_tokens=512, 
@@ -124,25 +104,20 @@ def inference_calling_idefics(image_path, question=""):
     )
     return response
 
-# Function to generate answers using the idefics model with system prompt and max_tokens
 def describe_image(image_path, question=""):
     try:
         answer = inference_calling_idefics(image_path, question)
         return answer
     except Exception as e:
         print(e)
-        answer = "Error while seeing the image."
-        return answer
+        return "Error while seeing the image."
 
-# Telegram bot client
 client = TelegramClient('bot', api_id, api_hash).start(bot_token=bot_token)
 
-# Get bot user ID for message filtering
 async def get_bot_id():
     me = await client.get_me()
     return me.id
 
-# Async function to get OpenAI completion
 async def get_completion(event, user_id, prompt):
     async with client.action(event.chat_id, 'typing'):
         await asyncio.sleep(3)
@@ -152,9 +127,8 @@ async def get_completion(event, user_id, prompt):
             *history,
             {"role": "user", "content": prompt},
         ]
-
         try:
-            response = openai_client.chat.completions.create(
+            completion = openai_client.chat.completions.create(
                 model=model,
                 messages=messages,
                 max_tokens=512,
@@ -162,17 +136,18 @@ async def get_completion(event, user_id, prompt):
                 top_p=1.0,
                 frequency_penalty=0.9,
                 presence_penalty=0.9,
+                stream=True
             )
-            message = response.choices[0].message.content
+            for chunk in completion:
+                if chunk.choices[0].delta.content is not None:
+                    message += chunk.choices[0].delta.content
         except Exception as e:
             message = f"Whoops!"
             print(e)
-
-        update_chat_history(user_id, "user", prompt)  # Update history
-        update_chat_history(user_id, "assistant", message)  # Update assistant's response
+        update_chat_history(user_id, "user", prompt)
+        update_chat_history(user_id, "assistant", message)
         return message
 
-# Telegram bot events
 @client.on(events.NewMessage(pattern='/start'))
 async def start(event):
     await event.respond("Hello!")
@@ -184,34 +159,27 @@ async def help(event):
 @client.on(events.NewMessage(pattern='/reset'))
 async def reset(event):
     user_id = event.chat_id
-    chat_collection.delete_one({"user_id": user_id})  # Reset MongoDB chat history for the user
+    chat_collection.delete_one({"user_id": user_id})
     if user_id in local_chat_history:
-        del local_chat_history[user_id]  # Remove from local cache if present
+        del local_chat_history[user_id]
     await event.respond("History has been reset.")
 
 @client.on(events.NewMessage)
 async def handle_message(event):
-    bot_id = await get_bot_id()  # Get bot ID to avoid responding to itself
-
+    async with client.action(event.chat_id, 'typing'):
+        await asyncio.sleep(1)
+    bot_id = await get_bot_id()
     try:
-        user_id = event.chat_id  # Use chat_id to distinguish between users
-
-        # Ignore messages from the bot itself
+        user_id = event.chat_id
         if event.sender_id == bot_id:
             return
-
         user_message = event.raw_text
-
         if event.photo:
-            # If an image is sent, describe the image
             photo = await event.download_media()
             image_description = describe_image(photo, user_message)
             user_message += f"\n\nI sent you an image. Content of the image: {image_description}"
-
-        # Ignore command messages to prevent double processing
         if user_message.startswith('/start') or user_message.startswith('/help') or user_message.startswith('/reset'):
             return
-
         response = await get_completion(event, user_id, user_message)
         await event.respond(response)
     except Exception as e:
@@ -232,14 +200,13 @@ def launch_gradio():
         """)
     demo.launch(show_api=False)
 
-# Keep-alive function to keep the bot running
 def keep_alive():
     ping_client = OpenAI(api_key=ping_key, base_url=api_url)
     while True:
         try:
             messages = [
-                {"role": "system", "content": "Be a helpful assistant."},
-                {"role": "user", "content": "Hello"}
+                {"role": "system", "content": "Repeat what i say."},
+                {"role": "user", "content": "Repeat: 'Ping success'"}
             ]
             request = ping_client.chat.completions.create(
                 model=model,
@@ -248,10 +215,10 @@ def keep_alive():
                 temperature=0.1,
                 top_p=0.1,
             )
-            print(request)
+            print(request.choices[0].message.content)
         except Exception as e:
             print(f"Keep-alive request failed: {e}")
-        time.sleep(1800)  # Ping every 30 minutes
+        time.sleep(1800)
 
 if __name__ == "__main__":
     threading.Thread(target=keep_alive).start()