Update README.md

README.md

@@ -22,17 +22,11 @@ tags:
 ## Code
 ```
 import os
-import ujson
-from typing import Optional
 
-os.environ['CUDA_VISIBLE_DEVICES'] = '2'
+os.environ['CUDA_VISIBLE_DEVICES'] = '6'
 
-import uvicorn
 import torch
 from transformers import AutoTokenizer, AutoModel
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
 
 from peft import PeftModel
 
@@ -47,27 +41,6 @@ tokenizer = AutoTokenizer.from_pretrained('THUDM/chatglm-6b',
                                           trust_remote_code=True)
 
 
-class ChatInfo(BaseModel):
-    owner: str
-    msg: str
-    unique_id: str
-
-
-class RatingInfo(BaseModel):
-    thumb_up: Optional[bool]
-    thumb_down: Optional[bool]
-    unique_id: str
-    idx: int
-
-
-app = FastAPI()
-app.add_middleware(CORSMiddleware,
-                   allow_origins=['*'],
-                   allow_credentials=True,
-                   allow_methods=['*'],
-                   allow_headers=['*'])
-
-
 def format_example(example: dict) -> dict:
     context = f'''Input: {example['input']}\n'''
 
@@ -88,30 +61,16 @@ def generate_response(data: dict):
                              temperature=0.9,
                              top_p=0.9)
 
-        raw_out_text = tokenizer.decode(out[0])
         true_out_text = tokenizer.decode(out[0][input_length:])
 
         answer = true_out_text.replace('\nEND', '').strip()
         return answer
 
 
-@app.post('/v1/chat')
-async def chat(ChatInfo: ChatInfo):
-    unique_id = ChatInfo.unique_id
-    # './dialogues'用于存储聊天数据
-    existing_files = os.listdir('./dialogues')
-
-    target_file = f'{unique_id}.json'
-    if target_file in existing_files:
-        with open(f'./dialogues/{unique_id}.json', 'r', encoding='utf-8') as f:
-            data: list = ujson.load(f)
-    else:
-        data = []
-    data.append({
-        'owner': ChatInfo.owner,
-        'msg': ChatInfo.msg,
-        'unique_id': ChatInfo.unique_id
-    })
+data = []
+while True:
+    seeker_msg = input('求助者：')
+    data.append({'owner': 'seeker', 'msg': seeker_msg})
     input_str = ''
     for item in data:
         if item['owner'] == 'seeker':
@@ -129,19 +88,9 @@ async def chat(ChatInfo: ChatInfo):
     wrapped_data = {'input': input_str}
 
     response = generate_response(data=wrapped_data)
-    supporter_msg = {
-        'owner': 'supporter',
-        'msg': response,
-        'unique_id': unique_id
-    }
+    print(f'支持者：{response}')
+    supporter_msg = {'owner': 'supporter', 'msg': response}
     data.append(supporter_msg)
-    with open(f'./dialogues/{unique_id}.json', 'w', encoding='utf-8') as f:
-        ujson.dump(data, f, ensure_ascii=False, indent=2)
-    return {'item': supporter_msg, 'responseCode': 200}
-
-
-if __name__ == '__main__':
-    uvicorn.run(app, host='0.0.0.0', port=8000)
 ```