Update app.py

app.py

@@ -1,62 +1,102 @@
 from pydantic import BaseModel
 from llama_cpp import Llama
 import os
-import gradio as gr
+import gradio as gr # Not suitable for production
 from dotenv import load_dotenv
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse
 import spaces
 import asyncio
 import random
-from io import BytesIO
-import requests
+from llama_cpp.tokenizers import LlamaTokenizer
+from peft import PeftModel, LoraConfig, get_peft_model
+import torch
+from multiprocessing import Process, Queue
+from google.cloud import storage
+import json
 
 app = FastAPI()
 load_dotenv()
 
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
+GOOGLE_CLOUD_BUCKET = os.getenv("GOOGLE_CLOUD_BUCKET")
+GOOGLE_CLOUD_CREDENTIALS = os.getenv("GOOGLE_CLOUD_CREDENTIALS")
+
+gcp_credentials = json.loads(GOOGLE_CLOUD_CREDENTIALS)
+storage_client = storage.Client.from_service_account_info(gcp_credentials)
+bucket = storage_client.bucket(GOOGLE_CLOUD_BUCKET)
 
 class ModelManager:
     def __init__(self):
-        self.params = {
-            "n_ctx": 2048,
-            "n_batch": 512,
-            "n_predict": 512,
-            "repeat_penalty": 1.1,
-            "n_threads": int(os.cpu_count() * 0.75),
-            "seed": -1,
-            "stop": ["</s>"],
-            "tokens": [],
-        }
-        self.unified_model = self.load_unified_model()
-
-    def load_unified_model(self):
-        model_configs = [
-            {
-                "repo_id": "TheBloke/Llama-2-7B-Chat-GGUF",
-                "filename": "llama-2-7b-chat.Q4_K_M.gguf",
-            },
-        ]
-
-        models = []
-        for config in model_configs:
-            with BytesIO() as model_data:
-                download_url = f"https://huggingface.co/{config['repo_id']}/resolve/main/{config['filename']}"
-                response = requests.get(download_url, headers={"Authorization": f"Bearer {HUGGINGFACE_TOKEN}"}, stream=True) 
-                
-                for chunk in response.iter_content(chunk_size=1024*1024):
-                    if chunk:
-                        model_data.write(chunk)
-                
-                model_data.seek(0)
-
-                model = Llama(model_path="", model_data=model_data.read(), **self.params)
-                models.append(model)
-
-                self.params["tokens"] = models[0].tokenize(b"Hello")
-
-        self.unified_model = models[0]
-        return self.unified_model
+        self.params = {"n_ctx": 2048, "n_batch": 512, "n_predict": 512, "repeat_penalty": 1.1, "n_threads": 1, "seed": -1, "stop": ["</s>"], "tokens": []}
+        self.tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf") #Load tokenizer from GCS for production
+        self.request_queue = Queue()
+        self.response_queue = Queue()
+        self.model = self.load_model_from_bucket("llama-2-7b-chat/llama-2-7b-chat.Q4_K_M.gguf")
+        self.start_processing_processes()
+
+    def load_model_from_bucket(self, bucket_path):
+        blob = bucket.blob(bucket_path)
+        try:
+            model = Llama(model_path=blob.download_as_string(), **self.params)
+            return model
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            return None
+
+    def save_model_to_bucket(self, model, bucket_path):
+        blob = bucket.blob(bucket_path)
+        try:
+            blob.upload_from_string(model.save_pretrained(), content_type='application/octet-stream')
+        except Exception as e:
+            print(f"Error saving model: {e}")
+
+    def train_model(self): #This function needs a complete overhaul for production use.  This is a placeholder.
+        config = LoraConfig(r=8, lora_alpha=32, target_modules=["q_proj", "v_proj"], lora_dropout=0.05, bias="none", task_type="CAUSAL_LM")
+        base_model_path = "llama-2-7b-chat/llama-2-7b-chat.Q4_K_M.gguf"
+        try:
+            base_model = self.load_model_from_bucket(base_model_path)
+            if base_model:
+                model = get_peft_model(base_model, config)
+                # Placeholder training data - needs a robust data loading mechanism
+                for batch in [{"question": ["a"], "answer":["b"]}, {"question":["c"], "answer":["d"]}]: 
+                    inputs = self.tokenizer(batch["question"], return_tensors="pt", padding=True, truncation=True)
+                    labels = self.tokenizer(batch["answer"], return_tensors="pt", padding=True, truncation=True)
+                    outputs = model(**inputs, labels=labels.input_ids)
+                    loss = outputs.loss
+                    loss.backward()
+                self.save_model_to_bucket(model, "llama_finetuned/llama_finetuned.gguf")
+                del model
+                del base_model
+        except Exception as e:
+            print(f"Error during training: {e}")
+
+
+    def generate_text(self, prompt):
+        if self.model:
+            inputs = self.tokenizer(prompt, return_tensors="pt")
+            outputs = self.model.generate(**inputs, max_new_tokens=100)
+            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            return generated_text
+        else:
+            return "Error loading model."
+
+    def start_processing_processes(self):
+        p = Process(target=self.process_requests)
+        p.start()
+
+    def process_requests(self):
+        while True:
+            request_data = self.request_queue.get()
+            if request_data is None:
+                break
+            inputs, top_p, top_k, temperature, max_tokens = request_data
+            try:
+                response = self.generate_text(inputs)
+                self.response_queue.put(response)
+            except Exception as e:
+                print(f"Error during inference: {e}")
+                self.response_queue.put("Error generating text.")
 
 model_manager = ModelManager()
 
@@ -65,22 +105,14 @@ class ChatRequest(BaseModel):
 
 @spaces.GPU()
 async def generate_streaming_response(inputs):
-    top_p = round(random.uniform(0.01, 1.00), 2)
-    top_k = random.randint(1, 100)
-    temperature = round(random.uniform(0.01, 2.00), 2)
-    max_tokens = model_manager.params["n_ctx"] - len(model_manager.unified_model.tokenize(inputs))
-
+    top_p = 0.9
+    top_k = 50
+    temperature = 0.7
+    max_tokens = model_manager.params["n_ctx"] - len(model_manager.tokenizer.encode(inputs))
+    model_manager.request_queue.put((inputs, top_p, top_k, temperature, max_tokens))
+    full_text = model_manager.response_queue.get()
     async def stream_response():
-        response = await asyncio.to_thread(model_manager.unified_model, inputs, top_p=top_p, top_k=top_k, temperature=temperature, max_tokens=max_tokens, **model_manager.params)
-        full_text = response['choices'][0]['text']
-
-        if len(full_text) > max_tokens:
-            chunks = [full_text[i:i + max_tokens] for i in range(0, len(full_text), max_tokens)]
-            for chunk in chunks:
-                yield chunk
-        else:
-            yield full_text
-
+        yield full_text
     return StreamingResponse(stream_response())
 
 async def process_message(message):
@@ -93,13 +125,7 @@ async def api_generate_multimodel(request: Request):
     message = data["message"]
     return await process_message(message)
 
-iface = gr.Interface(
-    fn=process_message,
-    inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."),
-    outputs=gr.Markdown(stream=True),
-    title="Unified Multi-Model API",
-    description="Enter a message to get responses from a unified model."
-)
+iface = gr.Interface(fn=process_message, inputs=gr.Textbox(lines=2, placeholder="Enter your message here..."), outputs=gr.Markdown(stream=True), title="Unified Multi-Model API", description="Enter a message to get responses from a unified model.") #gradio is not suitable for production
 
 if __name__ == "__main__":
     iface.launch()