Update app.py

app.py

@@ -1,85 +1,92 @@
-import gradio as gr
+from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
+from PIL import Image
+import requests
 import torch
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.document_loaders import TextLoader
-from langchain_community.vectorstores import FAISS
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.chains import RetrievalQA
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from langchain_huggingface import HuggingFacePipeline
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Load and process the document
-doc_loader = TextLoader("dataset.txt")
-docs = doc_loader.load()
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-split_docs = text_splitter.split_documents(docs)
-
-# Create embeddings and vector store
-embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
-vectordb = FAISS.from_documents(split_docs, embeddings)
-
-# Load model and tokenizer
-model_name = "01-ai/Yi-Coder-9B-Chat"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name, 
-    device_map="auto", 
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32
-)
+from threading import Thread
+import gradio as gr
+from gradio import FileData
+import time
+import spaces
+ckpt = "mrcuddle/llama3.2-11B-Vision_instruct-Coder"
+model = MllamaForConditionalGeneration.from_pretrained(ckpt,
+    torch_dtype=torch.bfloat16).to("cuda")
+processor = AutoProcessor.from_pretrained(ckpt)
 
-# Set up the QA pipeline
-qa_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=750,
-    pad_token_id=tokenizer.eos_token_id
-)
 
-llm = HuggingFacePipeline(pipeline=qa_pipeline)
+@spaces.GPU
+def bot_streaming(message, history, max_new_tokens=250):
+    
+    txt = message["text"]
+    ext_buffer = f"{txt}"
+    
+    messages= [] 
+    images = []
+    
 
-# Set up the retriever and QA chain
-retriever = vectordb.as_retriever(search_kwargs={"k": 5})
-qa_chain = RetrievalQA.from_chain_type(
-    retriever=retriever,
-    chain_type="stuff",
-    llm=llm,
-    return_source_documents=False
-)
+    for i, msg in enumerate(history): 
+        if isinstance(msg[0], tuple):
+            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
+            images.append(Image.open(msg[0][0]).convert("RGB"))
+        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
+            # messages are already handled
+            pass
+        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
+            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
+            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})
 
-def preprocess_query(query):
-    if "script" in query or "code" in query.lower():
-        return f"Write a CPSL script: {query}"
-    return query
+    # add current message
+    if len(message["files"]) == 1:
+        
+        if isinstance(message["files"][0], str): # examples
+            image = Image.open(message["files"][0]).convert("RGB")
+        else: # regular input
+            image = Image.open(message["files"][0]["path"]).convert("RGB")
+        images.append(image)
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
+    else:
+        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})
 
-def clean_response(response):
-    result = response.get("result", "")
-    if "Answer:" in result:
-        return result.split("Answer:")[1].strip()
-    return result.strip()
 
-def chatbot_response(user_input):
-    processed_query = preprocess_query(user_input)
-    raw_response = qa_chain.invoke({"query": processed_query})
-    return clean_response(raw_response)
+    texts = processor.apply_chat_template(messages, add_generation_prompt=True)
 
-# Gradio interface
-with gr.Blocks() as chat_interface:
-    gr.Markdown("# CPSL Chatbot")
-    chat_history = gr.Chatbot(type='messages')
-    user_input = gr.Textbox(label="Your Message:")
-    send_button = gr.Button("Send")
+    if images == []:
+        inputs = processor(text=texts, return_tensors="pt").to("cuda")
+    else:
+        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
+    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)
 
-    def interact(user_message, history):
-        bot_reply = chatbot_response(user_message)
-        history.append({"role": "user", "content": user_message})
-        history.append({"role": "assistant", "content": bot_reply})
-        return history, history
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens)
+    generated_text = ""
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    
+    for new_text in streamer:
+        buffer += new_text
+        generated_text_without_prompt = buffer
+        time.sleep(0.01)
+        yield buffer
 
-    send_button.click(interact, inputs=[user_input, chat_history], outputs=[chat_history, chat_history])
 
-# Launch the interface
-if __name__ == "__main__":
-    chat_interface.launch()
+demo = gr.ChatInterface(fn=bot_streaming, title="Multimodal Llama",examples=[
+    [{"text": "Replicate this webpage using Tyepescript and ChakraUI.", "files":["./examples/Untitled.png"]},
+    2000],
+],
+      textbox=gr.MultimodalTextbox(), 
+      additional_inputs = [gr.Slider(
+              minimum=10,
+              maximum=2500,
+              value=500,
+              step=10,
+              label="Maximum number of new tokens to generate",
+          )
+        ],
+      cache_examples=False,
+      description="Yes, this space can replicate (to the model's best ability) a webpage in your preferred language.",
+      stop_btn="Stop Generation", 
+      fill_height=True,
+    multimodal=True)
+    
+demo.launch(debug=True)