Mohammed-Altaf commited on
Commit
84ca46a
1 Parent(s): 1014b78

added new changes

Browse files
Files changed (3) hide show
  1. .gitignore +2 -0
  2. load_meta_data.py +51 -0
  3. requirements.txt +2 -0
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ __pycache__
2
+ venv
load_meta_data.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForCausalLM, AutoTokenizer
3
+ from pydantic import BaseModel
4
+ from fastapi import HTTPException
5
+
6
+ class UserQuery(BaseModel):
7
+ user_query: str
8
+
9
+
10
+ class ChatBot:
11
+ def __init__(self):
12
+ self.tokenizer = None
13
+ self.model = None
14
+
15
+ def load_from_hub(self,model_id: str):
16
+ self.tokenizer = AutoTokenizer.from_pretrained(model_id,
17
+ device_map = 'auto')
18
+ self.model = AutoModelForCausalLM.from_pretrained(model_id,
19
+ device_map='auto')
20
+
21
+ def get_response(self,text: UserQuery) -> str:
22
+ if not self.model or not self.tokenizer:
23
+ raise HTTPException(status_code=400, detail="Model is not loaded")
24
+
25
+ inputs = self.tokenizer(text,return_tensors='pt')
26
+ outputs = self.model.generate(**inputs,
27
+ max_new_tokens = 100,
28
+ )
29
+ response = self.tokenizer.decode(outputs[0],skip_special_tokens=True)
30
+ response = self.get_clean_response(response)
31
+ return response
32
+
33
+ def get_clean_response(self,response:list[str] | str) -> str:
34
+ if type(response) == list:
35
+ response = response[0].split("\n")
36
+ else:
37
+ response = response.split("\n")
38
+
39
+ ans = ''
40
+ cnt = 0 # to verify if we have seen Human before
41
+ for answer in response:
42
+ if answer.startswith("[|Human|]"): cnt += 1
43
+
44
+ elif answer.startswith('[|AI|]'):
45
+ answer = answer.split(' ')
46
+ ans += ' '.join(char for char in answer[1:])
47
+ ans += '\n'
48
+
49
+ elif cnt:
50
+ ans += answer + '\n'
51
+ return ans
requirements.txt ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ transformers
2
+ fastapi