Update app.py

app.py

@@ -1,41 +1,33 @@
 import gradio as gr
-import requests
-import os
-import json
+from openai import OpenAI
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer, util
 
-# 문장 임베딩 모델 로드
-model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+# OpenAI 클라이언트 초기화
+client = OpenAI(api_key=os.getenv("OPENAI"))  # 실제 API 키로 교체 필요
 
-# 데이터셋 로드
-datasets = [
-    ("all-processed", "all-processed"),
-    ("chatdoctor-icliniq", "chatdoctor-icliniq"),
-    ("chatdoctor_healthcaremagic", "chatdoctor_healthcaremagic"),
-]
+# Load sentence embedding model
+model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
 
-all_datasets = {}
-for dataset_name, config in datasets:
-    all_datasets[dataset_name] = load_dataset("lavita/medical-qa-datasets", config)
+# Load the PharmKG dataset
+pharmkg_dataset = load_dataset("vinven7/PharmKG")
 
 def find_most_similar_data(query):
     query_embedding = model.encode(query, convert_to_tensor=True)
     most_similar = None
     highest_similarity = -1
-    
-    for dataset_name, dataset in all_datasets.items():
-        for split in dataset.keys():
-            for item in dataset[split]:
-                if 'question' in item and 'answer' in item:
-                    item_text = f"질문: {item['question']} 답변: {item['answer']}"
-                    item_embedding = model.encode(item_text, convert_to_tensor=True)
-                    similarity = util.pytorch_cos_sim(query_embedding, item_embedding).item()
-                    
-                    if similarity > highest_similarity:
-                        highest_similarity = similarity
-                        most_similar = item_text
-    
+
+    for split in pharmkg_dataset.keys():
+        for item in pharmkg_dataset[split]:
+            if 'Input' in item and 'Output' in item:
+                item_text = f"Input: {item['Input']} Output: {item['Output']}"
+                item_embedding = model.encode(item_text, convert_to_tensor=True)
+                similarity = util.pytorch_cos_sim(query_embedding, item_embedding).item()
+
+                if similarity > highest_similarity:
+                    highest_similarity = similarity
+                    most_similar = item_text
+
     return most_similar
 
 def respond_with_prefix(message, history, max_tokens=10000, temperature=0.7, top_p=0.95):
@@ -86,51 +78,49 @@ def respond_with_prefix(message, history, max_tokens=10000, temperature=0.7, top
     7. 글의 전체가 아니라 챕터 마다 최소 1,000자 이상으로 세 챕터를 포함하면 3,000자 이상 작성해야 합니다.
     8. "#태그"를 10개 작성해주세요.
     """    
-
-    modified_message = system_prefix + message  # 사용자 메시지에 프리픽스 적용
+   
     
-    # 가장 유사한 데이터를 데이터셋에서 찾기
+    # Find the most similar data from PharmKG dataset
     similar_data = find_most_similar_data(message)
+    context = f"{system_prefix}\n\n{message}"
     if similar_data:
-        modified_message += "\n\n" + similar_data  # 유사한 데이터를 메시지에 추가
-    
-    data = {
-        "model": "jinjavis:latest",
-        "prompt": modified_message,
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p
-    }
-    
-    # API 요청
-    response = requests.post("http://hugpu.ai:7877/api/generate", json=data, stream=True)
-    
-    partial_message = ""
-    for line in response.iter_lines():
-        if line:
-            try:
-                result = json.loads(line)
-                if result.get("done", False):
-                    break
-                new_text = result.get('response', '')
-                partial_message += new_text
-                yield partial_message
-            except json.JSONDecodeError as e:
-                print(f"Failed to decode JSON: {e}")
-                yield "An error occurred while processing your request."
+        context += f"\n\nRelated Information: {similar_data}"
 
+    try:
+        response = client.chat.completions.create(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": system_prefix},
+                {"role": "user", "content": message}
+            ],
+            response_format={"type": "text"},
+            temperature=temperature,
+            max_tokens=max_tokens,
+            top_p=top_p,
+            frequency_penalty=0,
+            presence_penalty=0,
+            stream=True
+        )
 
+        partial_message = ""
+        for chunk in response:
+            if chunk.choices[0].delta.content:
+                partial_message += chunk.choices[0].delta.content
+                yield partial_message
 
-demo = gr.ChatInterface(
+    except Exception as e:
+        yield f"An error occurred: {str(e)}"
 
+demo = gr.ChatInterface(
     fn=respond_with_prefix,
     additional_inputs=[
-        gr.Slider(minimum=1, maximum=120000, value=4000, label="Max Tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, label="Top-P")  # Corrected comma placement
+        gr.Slider(minimum=1, maximum=4096, value=2048, label="Max Tokens"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=1.0, label="Top-P")
     ],
     theme="Nymbo/Nymbo_Theme"
 )
 
 if __name__ == "__main__":
     demo.queue(max_size=4).launch()
+